import gradio as gr
from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
from diffusers.utils import load_image
import torch

# Clear CUDA cache
torch.cuda.empty_cache()

# Define custom CSS to style the Gradio interface
css_style = """
.container {
    max-width: 3xl;
    padding: 2rem 1rem;
    margin: auto;
}
h1 {
    text-align: center;
    margin-bottom: 1.5rem;
    font-size: 1.875rem;
    font-weight: bold;
}
.form-textarea {
    width: 100%;
    margin-top: 0.25rem;
    height: 6rem;
}
.form-range {
    width: 100%;
}
.form-input {
    margin-top: 0.25rem;
    width: 9rem;
}
.form-checkbox {
    height: 1.25rem;
    width: 1.25rem;
    color: blue;
}
.button {
    background-color: blue;
    color: white;
    padding: 0.5rem 1rem;
    border-radius: 0.375rem;
    font-weight: bold;
}
.error {
    color: red;
    text-align: center;
}
"""
# Set environment variable for memory fragmentation
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

device = "cuda" if torch.cuda.is_available() else "cpu"

pipes = {
    "txt2img": AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to(device),
    "img2img": AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to(device)
}

if device == "cpu":
    pipes["txt2img"].enable_model_cpu_offload()
    pipes["img2img"].enable_model_cpu_offload()

def run(prompt, image):
    try:
        print(f"prompt={prompt}, image={image}")
        if image is None:
            return pipes["txt2img"](prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
        else:
            image = image.resize((512,512))
            print(f"img2img image={image}")
            return pipes["img2img"](prompt, image=image, num_inference_steps=2, strength=0.5, guidance_scale=0.0).images[0]
    except RuntimeError as e:
        if "CUDA out of memory" in str(e):
            print("CUDA out of memory. Trying to clear cache.")
            torch.cuda.empty_cache()
            # Consider additional fallback strategies here
        else:
            raise e

demo = gr.Interface(
    run,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Image(type="pil")
    ],
    outputs=gr.Image(width=512, height=512),
    live=True
)

demo.launch()