Update app.py

app.py

@@ -1,27 +1,27 @@
-# gradio_app.py
-
-import torch
-from diffusers import DiffusionPipeline
-import gradio as gr
-
-# Load model
-pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
-pipe.to("cuda")
-pipe.load_lora_weights("EliKet/train_text_to_img")
-
-# Generation function
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
-    return image
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your image prompt here..."),
-    outputs="image",
-    title="🐾 Lynx Text-to-Image Generator",
-    description="Type a prompt (e.g., 'A majestic lynx in a snowy forest') and get an AI-generated image using Stable Diffusion + LoRA."
-)
-
-# Launch app
-demo.launch()
+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
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4", torch_dtype=dtype
+)
+pipe.load_lora_weights("EliKet/train_text_to_img")
+pipe.to(device)
+
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+demo = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(placeholder="Enter your image prompt here..."),
+    outputs="image",
+    title="Text-to-Image Generator (Lynx)",
+    description="Type a prompt like 'a lynx in the snowy forest, ultra-detailed'."
+)
+
+demo.launch()