import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import gradio as gr
import spaces
import torch
from PIL import Image
from diffusers import StableDiffusionImg2ImgPipeline

pipe = None


@spaces.GPU
def generate(
    image,
    prompt,
    negative_prompt="",
    steps=10,
    strength=0.35,
):
    global pipe

    if image is None:
        raise gr.Error("Please upload an image")

    if pipe is None:
        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            torch_dtype=torch.float16,
            safety_checker=None,
            use_safetensors=True,
        )

        pipe.enable_attention_slicing()
        pipe.enable_vae_slicing()
        pipe.enable_model_cpu_offload()

    image = image.convert("RGB")
    image = image.resize((512, 512))

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image,
        num_inference_steps=int(steps),
        strength=float(strength),
        guidance_scale=7.5,
    )

    return result.images[0]


with gr.Blocks() as demo:
    gr.Markdown("# ImageGen2000")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                type="pil",
                label="Input Image"
            )

            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="photorealistic portrait"
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="blurry, low quality, distorted"
            )

            steps = gr.Slider(
                minimum=5,
                maximum=20,
                value=10,
                step=1,
                label="Steps"
            )

            strength = gr.Slider(
                minimum=0.2,
                maximum=0.7,
                value=0.35,
                step=0.05,
                label="Strength"
            )

            generate_btn = gr.Button("Generate")

        with gr.Column():
            output_image = gr.Image(label="Result")

    generate_btn.click(
        fn=generate,
        inputs=[
            input_image,
            prompt,
            negative_prompt,
            steps,
            strength,
        ],
        outputs=output_image,
    )

demo.launch()