import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from safetensors.torch import load_file

model_id = "runwayml/stable-diffusion-v1-5"
lora_path = "https://huggingface.co/codermert/model_malika/resolve/main/sarah-lora.safetensors"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

# LoRA dosyasını yükle
state_dict = load_file(lora_path)
pipe.unet.load_attn_procs(state_dict)

def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps
    ).images[0]
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Negative Prompt"),
        gr.Slider(minimum=1, maximum=20, step=0.5, label="Guidance Scale", value=7.5),
        gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=50)
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Stable Diffusion with LoRA",
    description="Generate images using Stable Diffusion v1.5 with a custom LoRA model."
)

iface.launch()