import gradio as gr
import torch
from diffusers import StableDiffusion3Pipeline

# Check if CUDA is available and set the device accordingly
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the Stable Diffusion 3.5 Large model
model_id = "stabilityai/stable-diffusion-3.5-large"
pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to(device)

# Define the image generation function
def generate_image(prompt, negative_prompt, width, height, guidance_scale, num_inference_steps, seed):
    generator = torch.manual_seed(seed) if seed else None
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator
    ).images[0]
    return image

# Set up the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion 3.5 Large Image Generator")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
            negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here")
            width = gr.Slider(label="Width", minimum=512, maximum=1024, step=64, value=512)
            height = gr.Slider(label="Height", minimum=512, maximum=1024, step=64, value=512)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=20.0, step=0.5, value=7.5)
            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=50)
            seed = gr.Number(label="Seed", value=42, precision=0)
            generate_button = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, width, height, guidance_scale, num_inference_steps, seed],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch()