import torch
from diffusers import StableDiffusionXLPipeline
import gradio as gr

# Load model
model_id = "OnomaAIResearch/Illustrious-XL-v1.1"

pipe = StableDiffusionXLPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
).to("cuda")

def generate_image(prompt, negative_prompt="", steps=30, guidance=7.5, seed=None):
    generator = torch.Generator(device="cuda").manual_seed(int(seed)) if seed else None
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(steps),
        guidance_scale=guidance,
        generator=generator,
    ).images[0]
    
    return image

with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Illustrious-XL Image Generator")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt")
            steps = gr.Slider(10, 50, value=30, label="Steps")
            guidance = gr.Slider(1.0, 15.0, value=7.5, label="Guidance Scale")
            seed = gr.Number(label="Seed", value=None)
            generate_btn = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(label="Result", height=512)

    generate_btn.click(
        generate_image,
        inputs=[prompt, negative_prompt, steps, guidance, seed],
        outputs=output_image
    )

demo.launch()