import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load SDXL Base 1.0 and apply your LoRA
base_model = "stabilityai/stable-diffusion-xl-base-1.0"
lora_model = "CoolKrishh/Comic-SDXL-LoRA"  # your LoRA model repo

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

pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch_dtype, use_safetensors=True)
pipe.load_lora_weights(lora_model)
pipe = pipe.to(device)

def infer(prompt):
    image = pipe(prompt=prompt, num_inference_steps=25, guidance_scale=7.5).images[0]
    return image

examples = [
    ["A superhero landing in a comic panel, cinematic lighting"],
    ["A futuristic city skyline drawn in comic style"],
    ["A knight fighting a dragon in anime-comic art"],
]

demo = gr.Interface(
    fn=infer,
    inputs=gr.Textbox(label="Prompt", placeholder="Describe your comic scene..."),
    outputs=gr.Image(label="Generated Comic Image"),
    examples=examples,
    title="Comic SDXL LoRA 🎨",
    description="Generate comic-style images using your fine-tuned SDXL LoRA model.",
)

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