Update app.py

app.py

@@ -1,125 +1,98 @@
 import os
-import requests
+import torch
 import gradio as gr
+from diffusers import DiffusionPipeline
 from PIL import Image, ImageDraw, ImageFont
-import io
-import time
-from concurrent.futures import ThreadPoolExecutor
 
-# ===== CONFIGURATION =====
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
-MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
-API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
-headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+# ===== FREE-TIER CONFIG =====
 WATERMARK_TEXT = "SelamGPT"
-MAX_RETRIES = 3
-TIMEOUT = 45
-EXECUTOR = ThreadPoolExecutor(max_workers=3)
+MODEL_NAME = "DeepFloyd/IF-II-L-v1.0"
+CACHE_DIR = "model_cache"  # For free tier storage limits
 
-# ===== OPTIMIZED WATERMARK FUNCTION =====
-def add_watermark(image_bytes):
+# ===== LIGHTWEIGHT MODEL LOAD =====
+pipe = None  # Lazy load to avoid cold start timeouts
+
+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
+        )
+        pipe.enable_model_cpu_offload()  # Critical for free-tier VRAM
+
+# ===== OPTIMIZED WATERMARK =====
+def add_watermark(image):
     try:
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         draw = ImageDraw.Draw(image)
-        
-        font_size = 24
-        try:
-            font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
-        except:
-            font = ImageFont.load_default(font_size)
-        
+        font = ImageFont.load_default(20)  # No external font needed
         text_width = draw.textlength(WATERMARK_TEXT, font=font)
-        x = image.width - text_width - 10
-        y = image.height - 34
-        
-        draw.text((x+1, y+1), WATERMARK_TEXT, font=font, fill=(0, 0, 0, 128))
-        draw.text((x, y), WATERMARK_TEXT, font=font, fill=(255, 255, 255))
-        
-        webp_buffer = io.BytesIO()
-        image.save(webp_buffer, format="WEBP", quality=85)
-        webp_buffer.seek(0)
-        return Image.open(webp_buffer)
-    except Exception as e:
-        print(f"Watermark error: {str(e)}")
-        return Image.open(io.BytesIO(image_bytes))
+        draw.text(
+            (image.width - text_width - 15, image.height - 30),
+            WATERMARK_TEXT,
+            font=font,
+            fill=(255, 255, 255)
+        )
+        return image
+    except Exception:
+        return image
 
-# ===== IMAGE GENERATION =====
+# ===== FREE-TIER GENERATION =====
 def generate_image(prompt):
     if not prompt.strip():
         return None, "⚠️ Please enter a prompt"
     
-    params = {
-        "height": 768,
-        "width": 768,
-        "num_inference_steps": 20,
-        "guidance_scale": 7.0,
-        "options": {"wait_for_model": False}
-    }
-    
-    def api_call():
-        return requests.post(
-            API_URL,
-            headers=headers,
-            json={
-                "inputs": prompt,
-                "parameters": params
-            },
-            timeout=TIMEOUT
+    try:
+        load_model()  # Lazy load only when needed
+        
+        # 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
         )
-    
-    for attempt in range(MAX_RETRIES):
-        try:
-            start_time = time.time()
-            response = EXECUTOR.submit(api_call).result()
-            
-            if response.status_code == 200:
-                gen_time = time.time() - start_time
-                return add_watermark(response.content), f"✔️ Generated in {gen_time:.1f}s"
-            elif response.status_code == 503:
-                time.sleep(5 * (attempt + 1))
-                continue
-            else:
-                return None, f"⚠️ API Error: {response.text[:200]}"
-        except requests.Timeout:
-            return None, f"⚠️ Timeout: Model took >{TIMEOUT}s"
-        except Exception as e:
-            return None, f"⚠️ Error: {str(e)[:200]}"
-    
-    return None, "⚠️ Failed after retries. Try again."
+        
+        return add_watermark(result.images[0]), "✔️ Generated (Free Tier)"
+    except torch.cuda.OutOfMemoryError:
+        return None, "⚠️ Out of VRAM - Try simpler prompt"
+    except Exception as e:
+        return None, f"⚠️ Error: {str(e)[:100]}"
 
-# ===== GRADIO INTERFACE =====
-with gr.Blocks(title="SelamGPT Image Generator") as demo:
+# ===== GRADIO UI =====
+with gr.Blocks(title="SelamGPT Pro") as demo:
     gr.Markdown("""
-    # 🎨 SelamGPT Image Generator
-    *Optimized for speed (WebP output @ 768px)*
+    # 🎨 SelamGPT (DeepFloyd IF-II-L)
+    *Optimized for Free Tier - 64px Base Resolution*
     """)
     
     with gr.Row():
-        with gr.Column(scale=3):
+        with gr.Column():
             prompt_input = gr.Textbox(
                 label="Describe your image",
-                placeholder="A futuristic Ethiopian city...",
+                placeholder="A traditional Ethiopian market...",
                 lines=3
             )
-            with gr.Row():
-                generate_btn = gr.Button("Generate Image", variant="primary")
-                clear_btn = gr.Button("Clear")
+            generate_btn = gr.Button("Generate", variant="primary")
             
             gr.Examples(
                 examples=[
-                    ["An ancient Aksumite warrior in cyberpunk armor"],
-                    ["Traditional Ethiopian coffee ceremony"],
-                    ["Habesha queen with jewelry"]
+                    ["Habesha cultural dress with intricate patterns, studio lighting"],
+                    ["Lalibela rock-hewn churches at golden hour"],
+                    ["Addis Ababa futuristic skyline, cyberpunk style"]
                 ],
                 inputs=prompt_input
             )
             
-        with gr.Column(scale=2):
+        with gr.Column():
             output_image = gr.Image(
-                label="Generated Image (WebP)",
+                label="Generated Image",
                 type="pil",
-                format="webp",
-                height=512
+                format="webp",  # Lightweight format
+                height=400
             )
             status_output = gr.Textbox(
                 label="Status",
@@ -129,20 +102,12 @@ with gr.Blocks(title="SelamGPT Image Generator") as demo:
     generate_btn.click(
         fn=generate_image,
         inputs=prompt_input,
-        outputs=[output_image, status_output],
-        queue=True
-    )
-    
-    clear_btn.click(
-        fn=lambda: [None, ""],
         outputs=[output_image, status_output]
     )
 
 if __name__ == "__main__":
-    # Version-compatible queue setup
-    try:
-        demo.queue(concurrency_count=3)  # New Gradio
-    except TypeError:
-        demo.queue(max_size=3)  # Old Gradio
-    
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        enable_queue=False  # Critical for free tier
+    )