Update app.py

app.py

@@ -7,27 +7,25 @@ from PIL import Image, ImageDraw, ImageFont
 # ===== FREE-TIER CONFIG =====
 WATERMARK_TEXT = "SelamGPT"
 MODEL_NAME = "DeepFloyd/IF-II-L-v1.0"
-CACHE_DIR = "model_cache"  # For free tier storage limits
 
-# ===== LIGHTWEIGHT MODEL LOAD =====
-pipe = None  # Lazy load to avoid cold start timeouts
+# Initialize pipeline (lazy load later)
+pipe = None
 
 def load_model():
     global pipe
     if pipe is None:
         pipe = DiffusionPipeline.from_pretrained(
             MODEL_NAME,
-            torch_dtype=torch.float16,  # 50% VRAM reduction
-            variant="fp16",
-            cache_dir=CACHE_DIR
+            torch_dtype=torch.float16,
+            variant="fp16"
         )
-        pipe.enable_model_cpu_offload()  # Critical for free-tier VRAM
+        pipe.to("cuda")
 
 # ===== OPTIMIZED WATERMARK =====
 def add_watermark(image):
     try:
         draw = ImageDraw.Draw(image)
-        font = ImageFont.load_default(20)  # No external font needed
+        font = ImageFont.load_default(20)
         text_width = draw.textlength(WATERMARK_TEXT, font=font)
         draw.text(
             (image.width - text_width - 15, image.height - 30),
@@ -39,50 +37,49 @@ def add_watermark(image):
     except Exception:
         return image
 
-# ===== FREE-TIER GENERATION =====
+# ===== GENERATION FUNCTION =====
 def generate_image(prompt):
     if not prompt.strip():
         return None, "⚠️ Please enter a prompt"
     
     try:
-        load_model()  # Lazy load only when needed
+        load_model()
         
-        # Free-tier optimized settings
         result = pipe(
             prompt=prompt,
             output_type="pil",
-            generator=torch.Generator().manual_seed(42),  # Consistent results
-            num_inference_steps=30,  # Reduced from default 50
-            guidance_scale=7.0  # Balanced creativity/quality
+            generator=torch.Generator(device="cuda").manual_seed(42),
+            num_inference_steps=30,
+            guidance_scale=7.0
         )
         
-        return add_watermark(result.images[0]), "✔️ Generated (Free Tier)"
+        return add_watermark(result.images[0]), "✔️ Generation successful"
     except torch.cuda.OutOfMemoryError:
-        return None, "⚠️ Out of VRAM - Try simpler prompt"
+        return None, "⚠️ Out of memory - Try a simpler prompt"
     except Exception as e:
-        return None, f"⚠️ Error: {str(e)[:100]}"
+        return None, f"⚠️ Error: {str(e)[:200]}"
 
-# ===== GRADIO UI =====
+# ===== GRADIO INTERFACE =====
 with gr.Blocks(title="SelamGPT Pro") as demo:
     gr.Markdown("""
     # 🎨 SelamGPT (DeepFloyd IF-II-L)
-    *Optimized for Free Tier - 64px Base Resolution*
+    *Free Tier Optimized - May take 2-3 minutes for first generation*
     """)
     
     with gr.Row():
         with gr.Column():
             prompt_input = gr.Textbox(
                 label="Describe your image",
-                placeholder="A traditional Ethiopian market...",
+                placeholder="A traditional Ethiopian market scene...",
                 lines=3
             )
             generate_btn = gr.Button("Generate", variant="primary")
             
             gr.Examples(
                 examples=[
-                    ["Habesha cultural dress with intricate patterns, studio lighting"],
-                    ["Lalibela rock-hewn churches at golden hour"],
-                    ["Addis Ababa futuristic skyline, cyberpunk style"]
+                    ["Habesha cultural dress with gold embroidery, studio lighting"],
+                    ["Lalibela churches at sunrise, foggy morning"],
+                    ["Futuristic Addis Ababa with Ethiopian architecture"]
                 ],
                 inputs=prompt_input
             )
@@ -91,7 +88,7 @@ with gr.Blocks(title="SelamGPT Pro") as demo:
             output_image = gr.Image(
                 label="Generated Image",
                 type="pil",
-                format="webp",  # Lightweight format
+                format="webp",
                 height=400
             )
             status_output = gr.Textbox(
@@ -108,6 +105,5 @@ with gr.Blocks(title="SelamGPT Pro") as demo:
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        server_port=7860,
-        enable_queue=False  # Critical for free tier
+        server_port=7860
     )