Update app.py

app.py

@@ -4,107 +4,175 @@ import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
 import io
 import time
+from concurrent.futures import ThreadPoolExecutor
 
-# Configuration
+# ===== CONFIGURATION =====
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
-MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"  # High-quality model
+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}"}
 WATERMARK_TEXT = "SelamGPT"
 MAX_RETRIES = 3
+TIMEOUT = 45  # Increased timeout for larger images
+EXECUTOR = ThreadPoolExecutor(max_workers=2)  # Handle concurrent requests
 
+# ===== WATERMARK FUNCTION =====
 def add_watermark(image_bytes):
-    """Add watermark to generated image"""
+    """Add professional watermark with fallback fonts"""
     try:
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         draw = ImageDraw.Draw(image)
         
-        # Use default font (no need for external files)
-        try:
-            font = ImageFont.truetype("arialbd.ttf", 40)
-        except:
-            font = ImageFont.load_default(size=40)  # Fallback
+        # Try multiple font options
+        font = None
+        for font_path in [
+            "Roboto-Bold.ttf",  # Our Dockerfile installs this
+            "DejaVuSans-Bold.ttf",
+            "FreeSansBold.ttf",
+            None  # Final fallback to default
+        ]:
+            try:
+                size = min(image.width // 15, 40)  # Responsive sizing
+                font = ImageFont.truetype(font_path, size) if font_path else ImageFont.load_default(size)
+                break
+            except:
+                continue
         
-        # Calculate text position (bottom-right corner)
+        # Calculate dynamic position
         bbox = draw.textbbox((0, 0), WATERMARK_TEXT, font=font)
-        text_width = bbox[2] - bbox[0]
-        text_height = bbox[3] - bbox[1]
-        margin = 20
-        position = (image.width - text_width - margin, image.height - text_height - margin)
+        text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+        margin = image.width // 50
+        position = (image.width - text_w - margin, image.height - text_h - margin)
+        
+        # Draw with outline effect
+        for offset in [(-1,-1), (1,1)]:  # Shadow positions
+            draw.text(
+                (position[0]+offset[0], position[1]+offset[1]),
+                WATERMARK_TEXT,
+                font=font,
+                fill=(0, 0, 0, 180)  # Semi-transparent black
         
-        # Draw watermark with semi-transparent white text and black outline
         draw.text(
             position,
             WATERMARK_TEXT,
             font=font,
-            fill=(255, 255, 255, 128),
-            stroke_width=2,
-            stroke_fill=(0, 0, 0, 128)
-        )
+            fill=(255, 255, 255, 200))  # Semi-transparent white
         
         return image
     except Exception as e:
         print(f"Watermark error: {str(e)}")
-        return Image.open(io.BytesIO(image_bytes))  # Return original if watermark fails
+        return Image.open(io.BytesIO(image_bytes))  # Fallback to original
 
+# ===== IMAGE GENERATION =====
 def generate_image(prompt):
-    """Generate image with retry logic"""
+    """Generate image with robust error handling"""
     if not prompt.strip():
-        return "Error: Please enter a valid prompt"
+        return None, "⚠️ Please enter a prompt"
+    
+    def api_call():
+        return requests.post(
+            API_URL,
+            headers=headers,
+            json={
+                "inputs": prompt,
+                "parameters": {
+                    "height": 768,
+                    "width": 768,
+                    "num_inference_steps": 25
+                },
+                "options": {"wait_for_model": True}
+            },
+            timeout=TIMEOUT
+        )
     
     for attempt in range(MAX_RETRIES):
         try:
-            response = requests.post(
-                API_URL,
-                headers=headers,
-                json={"inputs": prompt, "options": {"wait_for_model": True}},
-                timeout=30
-            )
+            future = EXECUTOR.submit(api_call)
+            response = future.result()
             
             if response.status_code == 200:
-                return add_watermark(response.content)
-            elif response.status_code == 503:  # Model loading
-                time.sleep(10 * (attempt + 1))  # Exponential backoff
+                return add_watermark(response.content), "✔️ Generation successful"
+            elif response.status_code == 503:
+                wait_time = (attempt + 1) * 10
+                print(f"Model loading, waiting {wait_time}s...")
+                time.sleep(wait_time)
                 continue
             else:
-                return f"API Error: {response.text}"
+                return None, f"⚠️ API Error: {response.text[:200]}"
         except requests.Timeout:
-            return "Error: Request timed out (30s)"
+            return None, "⚠️ Timeout: Model took too long to respond"
         except Exception as e:
-            return f"Unexpected error: {str(e)}"
+            return None, f"⚠️ Unexpected error: {str(e)[:200]}"
     
-    return "Failed after multiple attempts. Try again later."
+    return None, "⚠️ Failed after multiple attempts. Please try again later."
+
+# ===== GRADIO INTERFACE =====
+theme = gr.themes.Default(
+    primary_hue="emerald",
+    secondary_hue="amber",
+    font=[gr.themes.GoogleFont("Poppins"), "Arial", "sans-serif"]
+)
 
-# Gradio Interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
     gr.Markdown("""
     # 🎨 SelamGPT Image Generator
-    *Powered by Stable Diffusion XL with built-in watermark*
+    *Generate watermarked images with Stable Diffusion XL*
     """)
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=3):
             prompt_input = gr.Textbox(
                 label="Describe your image",
-                placeholder="A futuristic city at sunset...",
-                lines=3
+                placeholder="A futuristic Ethiopian city with flying cars...",
+                lines=3,
+                max_lines=5,
+                elem_id="prompt-box"
             )
-            generate_btn = gr.Button("Generate", variant="primary")
-            examples = gr.Examples(
-                examples=["A cute robot reading a book", "Ethiopian landscape in oil painting style"],
-                inputs=prompt_input
+            with gr.Row():
+                generate_btn = gr.Button("Generate Image", variant="primary")
+                clear_btn = gr.Button("Clear")
+            
+            gr.Examples(
+                examples=[
+                    ["An ancient Aksumite warrior in cyberpunk armor"],
+                    ["Traditional Ethiopian coffee ceremony in space"],
+                    ["Hyper-realistic portrait of a Habesha woman with neon tribal markings"]
+                ],
+                inputs=prompt_input,
+                label="Example Prompts"
+            )
+            
+        with gr.Column(scale=2):
+            output_image = gr.Image(
+                label="Generated Image",
+                height=512,
+                elem_id="output-image"
+            )
+            status_output = gr.Textbox(
+                label="Status",
+                interactive=False,
+                elem_id="status-box"
             )
-        
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image", height=512)
-            error_output = gr.Textbox(label="Status", visible=False)
     
+    # Event handlers
     generate_btn.click(
         fn=generate_image,
         inputs=prompt_input,
-        outputs=[output_image, error_output],
-        show_progress="full"
+        outputs=[output_image, status_output],
+        queue=True,
+        show_progress="minimal"
+    )
+    
+    clear_btn.click(
+        fn=lambda: [None, ""],
+        outputs=[output_image, status_output]
     )
 
+# ===== DEPLOYMENT CONFIG =====
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue(concurrency_count=2, api_open=False)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        favicon_path="./favicon.ico"  # Optional
+    )