import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

model_id = "stabilityai/stable-diffusion-2-1"
lora_id = "cocktailpeanut/itojunjilora"

# Load Stable Diffusion
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")

# Load LoRA
pipe.load_lora_weights(lora_id)
pipe.fuse_lora()

# Function to generate image
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="LoRA Stable Diffusion",
    description="Nhập mô tả hình ảnh để tạo ra với LoRA."
)

interface.launch()