import gradio as gr
from PIL import Image
from loguru import logger
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers import AutoPipelineForText2Image


# Init pipelines
pipe_t2i = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe_t2i.to("cuda")

pipe_i2i = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe_i2i.to("cuda")


def generate_image(text_input, prev_image, strength):
    if prev_image:
        logger.info(f"iterating for prompt: {text_input} with strength {strength}")
        image = pipe_i2i(
            prompt=text_input,
            num_inference_steps=4,
            guidance_scale=0.0,
            strength=strength,
            image=prev_image
        ).images[0]
    else:
        logger.info(f"generating for prompt: {text_input}")
        image = pipe_t2i(text_input, num_inference_steps=4, guidance_scale=0.0).images[0]
    return image


# Create left block (Chatbot)
with gr.Blocks() as demo:
    with gr.Row() as row:
        with gr.Column():
            # create left block (Text input)
            chatbot = gr.Chatbot(label="History")
            msg = gr.Textbox(label="Input prompt and press Enter")
            strength_slider = gr.Slider(minimum=0, maximum=1, value=0.75, label="Creativity of image iterations")
            regen = gr.Button("Regenerate for the last prompt")
            clear = gr.Button("Clear")

            def respond(message, chat_history, gen_image, strength_slider, prev_gen_image):
                user_input = [chat_history[i][0] for i in range(len(chat_history))]
                user_input.append(message)
                bot_message = "Generating image for: " + " ".join(user_input)
                chat_history.append((message, bot_message))
                if gen_image:
                    prev_gen_image = gen_image
                else:
                    prev_gen_image = None
                gen_image = generate_image(" ".join(user_input), gen_image, strength_slider)
                return "", chat_history, gen_image, prev_gen_image

            def regenerate(message, chat_history, gen_image, strength_slider):
                # Regenerate with the previous image
                user_input = [chat_history[i][0] for i in range(len(chat_history))]
                # user_input.append(message)
                bot_message = "Regenerating image for: " + " ".join(user_input)
                chat_history.append((message, bot_message))
                gen_image = generate_image(" ".join(user_input), gen_image, strength_slider)
                return "", chat_history, gen_image

        with gr.Column():
            # Create right block (Image Generator)
            image_generator = gr.Image(label="Generated image", type="pil", interactive=False, show_download_button=True)
            # use this to keep track of previous image generation but don't display it
            image_generator_prev = gr.Image(label="Previous image", type="pil", interactive=False)

    msg.submit(respond, [msg, chatbot, image_generator, strength_slider, image_generator_prev], [msg, chatbot, image_generator, image_generator_prev])
    regen.click(regenerate, [msg, chatbot, image_generator_prev, strength_slider], [msg, chatbot, image_generator])
    clear.click(lambda: [None, None, None, 0.75], None, [chatbot, image_generator, image_generator_prev, strength_slider], queue=False)


# Launch the Gradio app
demo.launch(share=True)