Update app.py

app.py

@@ -2,26 +2,33 @@ import torch
 from diffusers import DiffusionPipeline
 import gradio as gr
 
-# Load model with float16 for GPU or float32 for CPU
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+# Detect device
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
+# Load pipeline
 pipe = DiffusionPipeline.from_pretrained(
-    "CompVis/stable-diffusion-v1-4", torch_dtype=dtype
+    "CompVis/stable-diffusion-v1-4",
+    torch_dtype=dtype
 )
-pipe.load_lora_weights("EliKet/train_text_to_img")
 pipe.to(device)
 
+# Load LoRA weights (requires `peft` installed)
+pipe.load_lora_weights("EliKet/train_text_to_img")
+
+# Inference function
 def generate_image(prompt):
     image = pipe(prompt).images[0]
     return image
 
+# Gradio Interface
 demo = gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(placeholder="Enter your image prompt here..."),
+    inputs=gr.Textbox(lines=2, placeholder="Describe the image you want..."),
     outputs="image",
-    title="Text-to-Image Generator (Lynx)",
-    description="Type a prompt like 'a lynx in the snowy forest, ultra-detailed'."
+    title="🖼️ LoRA Text-to-Image Generator",
+    description="Enter a prompt to generate an image using Stable Diffusion with LoRA (EliKet/train_text_to_img)."
 )
 
+# Launch app
 demo.launch()