Update app.py

app.py

@@ -1,33 +1,30 @@
 import torch
-import gradio as gr
 from diffusers import StableDiffusionPipeline
+import gradio as gr
 
-# Load model từ Hugging Face
-model_id = "cocktailpeanut/itojunjidiffusion"
-pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+model_id = "stabilityai/stable-diffusion-2-1"
+lora_id = "cocktailpeanut/itojunjilora"
 
-# Hàm sinh ảnh từ prompt
-def generate_image(prompt, steps, guidance_scale):
-    image = pipeline(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
-    return image
+# Load Stable Diffusion
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe.to("cuda")
 
-# Tạo giao diện với Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# 🌀 Ito Junji Diffusion - Horror Art Generator")
-    gr.Markdown("Nhập mô tả để tạo ảnh theo phong cách kinh dị của Junji Ito!")
+# Load LoRA
+pipe.load_lora_weights(lora_id)
+pipe.fuse_lora()
 
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt", value="A terrifying horror manga scene")
-            steps = gr.Slider(10, 50, value=25, step=1, label="Inference Steps")
-            guidance_scale = gr.Slider(1, 10, value=7.5, step=0.1, label="Guidance Scale")
-            btn = gr.Button("Generate Image")
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image")
+# Function to generate image
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
 
-    btn.click(generate_image, inputs=[prompt, steps, guidance_scale], outputs=output_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."
+)
 
-# Chạy app
-if __name__ == "__main__":
-    demo.launch()
+interface.launch()