Update app.py

app.py

@@ -8,28 +8,21 @@ from concurrent.futures import ThreadPoolExecutor
 
 # ===== CONFIGURATION =====
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
-MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"  # Using SDXL
+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 = 60  # Increased for SDXL's longer processing
+TIMEOUT = 60
 EXECUTOR = ThreadPoolExecutor(max_workers=2)
 
 # ===== WATERMARK FUNCTION =====
 def add_watermark(image_bytes):
-    """Convert to PNG with medium quality before watermarking"""
+    """Add watermark with optimized PNG output"""
     try:
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        draw = ImageDraw.Draw(image)
         
-        # Save as medium-quality PNG to buffer
-        png_buffer = io.BytesIO()
-        image.save(png_buffer, format="PNG", optimize=True, quality=85)  # Medium quality
-        png_buffer.seek(0)
-        
-        # Add watermark to the PNG
-        watermarked_image = Image.open(png_buffer)
-        draw = ImageDraw.Draw(watermarked_image)
         font_size = 24
         try:
             font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
@@ -37,25 +30,23 @@ def add_watermark(image_bytes):
             font = ImageFont.load_default(font_size)
         
         text_width = draw.textlength(WATERMARK_TEXT, font=font)
-        x = watermarked_image.width - text_width - 10
-        y = watermarked_image.height - 34
+        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))
         
-        # Return as PNG bytes
-        final_buffer = io.BytesIO()
-        watermarked_image.save(final_buffer, format="PNG", optimize=True, quality=85)
-        final_buffer.seek(0)
-        return Image.open(final_buffer)
-        
+        # Convert to optimized PNG
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG', optimize=True, quality=85)
+        img_byte_arr.seek(0)
+        return Image.open(img_byte_arr)
     except Exception as e:
         print(f"Watermark error: {str(e)}")
         return Image.open(io.BytesIO(image_bytes))
 
-# ===== IMAGE GENERATION (SDXL-OPTIMIZED) =====
+# ===== IMAGE GENERATION =====
 def generate_image(prompt):
-    """Generate image with SDXL-specific parameters"""
     if not prompt.strip():
         return None, "⚠️ Please enter a prompt"
     
@@ -66,10 +57,9 @@ def generate_image(prompt):
             json={
                 "inputs": prompt,
                 "parameters": {
-                    "height": 1024,  # SDXL's native resolution
+                    "height": 1024,
                     "width": 1024,
-                    "num_inference_steps": 30,  # Better quality than 25
-                    "guidance_scale": 7.5  # SDXL's optimal value
+                    "num_inference_steps": 30
                 },
                 "options": {"wait_for_model": True}
             },
@@ -84,7 +74,7 @@ def generate_image(prompt):
             if response.status_code == 200:
                 return add_watermark(response.content), "✔️ Generation successful"
             elif response.status_code == 503:
-                wait_time = (attempt + 1) * 15  # Longer wait for SDXL
+                wait_time = (attempt + 1) * 15
                 print(f"Model loading, waiting {wait_time}s...")
                 time.sleep(wait_time)
                 continue
@@ -97,20 +87,18 @@ def generate_image(prompt):
     
     return None, "⚠️ Failed after multiple attempts. Please try later."
 
-# ===== GRADIO INTERFACE =====
-with gr.Blocks() as demo:
-    
-    output_image = gr.Image(
-        label="Generated Image",
-        type="pil",  # Force PIL/PNG output
-        format="png",  # Explicit PNG format
-        height=512
-    )
+# ===== GRADIO THEME =====
+theme = gr.themes.Default(
+    primary_hue="emerald",
+    secondary_hue="amber",
+    font=[gr.themes.GoogleFont("Poppins"), "Arial", "sans-serif"]
+)
 
+# ===== GRADIO INTERFACE =====
 with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
     gr.Markdown("""
     # 🎨 SelamGPT Image Generator
-    *Now powered by Stable Diffusion XL (1024x1024 resolution)*
+    *Powered by Stable Diffusion XL (1024x1024 PNG output)*
     """)
     
     with gr.Row():
@@ -119,8 +107,7 @@ with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
                 label="Describe your image",
                 placeholder="A futuristic Ethiopian city with flying cars...",
                 lines=3,
-                max_lines=5,
-                elem_id="prompt-box"
+                max_lines=5
             )
             with gr.Row():
                 generate_btn = gr.Button("Generate Image", variant="primary")
@@ -129,31 +116,29 @@ with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
             gr.Examples(
                 examples=[
                     ["An ancient Aksumite warrior in cyberpunk armor, 4k detailed"],
-                    ["Traditional Ethiopian coffee ceremony in zero gravity, photorealistic"],
-                    ["Portrait of a Habesha queen with golden jewelry, studio lighting"]
+                    ["Traditional Ethiopian coffee ceremony in zero gravity"],
+                    ["Portrait of a Habesha queen with golden jewelry"]
                 ],
-                inputs=prompt_input,
-                label="Try these SDXL-optimized prompts:"
+                inputs=prompt_input
             )
             
         with gr.Column(scale=2):
             output_image = gr.Image(
-                label="Generated Image (1024x1024)",
-                height=512,
-                elem_id="output-image"
+                label="Generated Image",
+                type="pil",
+                format="png",
+                height=512
             )
             status_output = gr.Textbox(
                 label="Status",
-                interactive=False,
-                elem_id="status-box"
+                interactive=False
             )
     
     generate_btn.click(
         fn=generate_image,
         inputs=prompt_input,
         outputs=[output_image, status_output],
-        queue=True,
-        show_progress="minimal"
+        queue=True
     )
     
     clear_btn.click(