import gradio as gr
import replicate
from openai import OpenAI
from PIL import Image
import requests
from io import BytesIO
import os

OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
REPLICATE_API_TOKEN = os.getenv('REPLICATE_API_TOKEN')


def generate_image_openai(prompt):

    client = OpenAI()
    response = client.images.generate(
        model="dall-e-3",
        prompt=prompt,
        size="1024x1024",
        n=1,
    )

    return response.data[0].url, response.data[0].revised_prompt

def style_transfer(input_image_path, style_image_path, revised_prompt):
    input = {
        # "image": open(input_image_path, "rb"),
        "image": input_image_path,
        "image_style": open(style_image_path, "rb"),
        "style_strength": 0.4,
        "structure_strength":1.2,
        "prompt": " natural light, natural bright colors, low quality, candid, grainy, instagram photo, phone camera, candid, blurry home video, high iso noisy" ,
        "seed": 42,
    }
    output = replicate.run(
        "prakharsaxena24/2d-to-real-style:c0e1e612a11a13d1d57a6d647e7665ad850bc73715337c1f499bb7b52404c35a",
        input=input
    )
    return output[0]



def infer(text,title):

    prompt = f"""Please create a simple suitable image to accompany the following text as part of an article with the title "{title}". The objects in the image must have realistic proportions. Text: "{text}" 
     Please make sure not to include text in the image."""
    image_url_openai, revised_prompt = generate_image_openai(prompt)
    style_image_url = style_transfer(image_url_openai, f'./Style.png',revised_prompt)
    response = requests.get(style_image_url)
    img = Image.open(BytesIO(response.content))

    return img
    
    
    
    


examples = [
    # "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    # "An astronaut riding a green horse",
    # "A delicious ceviche cheesecake slice",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""



with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # Text-to-Image and style transfer.
        """)
        
        with gr.Row():
            
            text = gr.Text(
                label="Text",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            title = gr.Text(
                label="Title",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

        # with gr.Accordion("Advanced Settings", open=False):
            
        #     negative_prompt = gr.Text(
        #         label="Negative prompt",
        #         max_lines=1,
        #         placeholder="Enter a negative prompt",
        #         visible=False,
        #     )
            
        #     seed = gr.Slider(
        #         label="Seed",
        #         minimum=0,
        #         maximum=100000,
        #         step=1,
        #         value=0,
        #     )
            
        #     randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
        #     with gr.Row():
                
        #         width = gr.Slider(
        #             label="Width",
        #             minimum=256,
        #             maximum=MAX_IMAGE_SIZE,
        #             step=32,
        #             value=512,
        #         )
                
        #         height = gr.Slider(
        #             label="Height",
        #             minimum=256,
        #             maximum=MAX_IMAGE_SIZE,
        #             step=32,
        #             value=512,
        #         )
            
        #     with gr.Row():
                
        #         guidance_scale = gr.Slider(
        #             label="Guidance scale",
        #             minimum=0.0,
        #             maximum=10.0,
        #             step=0.1,
        #             value=0.0,
        #         )
                
        #         num_inference_steps = gr.Slider(
        #             label="Number of inference steps",
        #             minimum=1,
        #             maximum=12,
        #             step=1,
        #             value=2,
        #         )
        
        # gr.Examples(
        #     examples = examples,
        #     inputs = [prompt]
        # )

    run_button.click(
        fn = infer,
        inputs = [text, title],
        outputs = [result]
    )

demo.queue().launch()