import gradio as gr
import torch
import spaces
from diffusers import DiffusionPipeline

pipe = None

@spaces.GPU(duration=180)
def generate(prompt, negative_prompt, steps, guidance, seed):
    global pipe
    if pipe is None:
        pipe = DiffusionPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True
        ).to("cuda")
        pipe.load_lora_weights("dc-design/kira-sdxl-lora", weight_name="kira_lora.safetensors")
        print("Model loaded!")
    
    generator = torch.Generator("cuda").manual_seed(int(seed))
    image = pipe(
        "a photo of ohwx " + prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=768,
        height=1024,
        generator=generator
    ).images[0]
    return image

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", value="woman at a coffee shop, warm lighting, photorealistic"),
        gr.Textbox(label="Negative", value="ugly, deformed, blurry, cgi, cartoon"),
        gr.Slider(10, 30, value=20, step=1, label="Steps"),
        gr.Slider(1, 15, value=7.0, step=0.5, label="Guidance"),
        gr.Number(value=42, label="Seed"),
    ],
    outputs=gr.Image(label="Result"),
    title="Kira AI Influencer Generator",
)

demo.launch(show_error=True)