from diffusers import StableDiffusionXLPipeline
import torch
from gradio import Interface, Image, Dropdown, Slider
import gradio as gr
import spaces

model_id = "RunDiffusion/Juggernaut-X-v10"
pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

@spaces.GPU()
def text_to_image(prompt, negative_prompt, steps, guidance_scale, add_4k_masterpiece, width, height, progress=gr.Progress(track_tqdm=True)):
    if add_4k_masterpiece:
        prompt += ", 4k, (masterpiece)"
    image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=guidance_scale, width=width, height=height).images[0]
    return image

duplicate_button = gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")

def set_aspect_ratio(ratio):
    ratios = {
        "1:1": (1024, 1024),
        "4:3": (1366, 1024),
        "16:9": (1920, 1080),
        "3:2": (2048, 1365),
        "5:4": (1280, 1024)
    }
    return ratios[ratio]

def update_aspect_ratio(ratio, width_slider, height_slider):
    width, height = set_aspect_ratio(ratio)
    width_slider.update(value=width)
    height_slider.update(value=height)

with gr.Blocks() as gradio_interface:
    gr.Markdown("## Stable Diffusion XL Interface")
    
    dark_mode = gr.Checkbox(label="Dark Mode")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Enter your prompt here...")
        negative_prompt = gr.Textbox(label="Negative Prompt", lines=2, placeholder="What to exclude from the image...")
    
    with gr.Row():
        steps = gr.Slider(minimum=1, maximum=65, value=50, label="Steps", step=1)
        guidance_scale = gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale", step=0.1)
    
    add_4k_masterpiece = gr.Checkbox(label="Add recommended prompt items (4k, masterpiece)", value=False)
    
    with gr.Row():
        width = gr.Slider(minimum=0, maximum=2048, value=1024, label="Width", step=1)
        height = gr.Slider(minimum=0, maximum=2048, value=1024, label="Height", step=1)
    
    aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["1:1", "4:3", "16:9", "3:2", "5:4"])
    
    aspect_ratio.change(update_aspect_ratio, inputs=[aspect_ratio, width, height], outputs=[width, height])
    
    generate_button = gr.Button("Generate Image")
    output_image = gr.Image(type="pil", show_download_button=True)

    generate_button.click(
        fn=text_to_image, 
        inputs=[prompt, negative_prompt, steps, guidance_scale, add_4k_masterpiece, width, height],
        outputs=output_image
    )

    if dark_mode:
        gradio_interface.theme = gr.themes.Dark()

gradio_interface.launch()