Update app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
 import io
@@ -11,25 +10,21 @@ WATERMARK_TEXT = "SelamGPT"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 TORCH_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
-# ===== MODEL LOADING =====
-# Global variable for model caching (alternative to @gr.Cache)
-pipe = None
-
+# ===== MODEL LOADING WITH GRADIO CACHE =====
+@gr.Cache()  # Now works in Gradio 5.x
 def load_model():
-    global pipe
-    if pipe is None:
-        pipe = DiffusionPipeline.from_pretrained(
-            MODEL_NAME,
-            torch_dtype=TORCH_DTYPE
-        ).to(DEVICE)
-        
-        # Optimizations
-        if DEVICE == "cuda":
-            try:
-                pipe.enable_xformers_memory_efficient_attention()
-            except:
-                print("Xformers not available, using default attention")
-            pipe.enable_attention_slicing()
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=TORCH_DTYPE
+    ).to(DEVICE)
+    
+    # Optimizations
+    if DEVICE == "cuda":
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+        except:
+            print("Xformers not available, using default attention")
+        pipe.enable_attention_slicing()
     
     return pipe
 
@@ -39,23 +34,25 @@ def add_watermark(image):
     try:
         draw = ImageDraw.Draw(image)
         
-        font_size = 24
+        font_size = max(24, int(image.width * 0.03))  # Dynamic font sizing
         try:
             font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
         except:
             font = ImageFont.load_default(font_size)
         
         text_width = draw.textlength(WATERMARK_TEXT, font=font)
-        x = image.width - text_width - 10
-        y = image.height - 34
+        margin = image.width * 0.02  # Dynamic margin
+        x = image.width - text_width - margin
+        y = image.height - (font_size * 1.5)
         
-        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))
+        # Shadow effect
+        draw.text((x+2, y+2), WATERMARK_TEXT, font=font, fill=(0, 0, 0, 150))
+        # Main text
+        draw.text((x, y), WATERMARK_TEXT, font=font, fill=(255, 215, 0))  # Gold color
         
-        # Convert to optimized PNG
+        # Optimized PNG output
         img_byte_arr = io.BytesIO()
-        image.save(img_byte_arr, format='PNG', optimize=True, quality=85)
-        img_byte_arr.seek(0)
+        image.save(img_byte_arr, format='PNG', optimize=True)
         return Image.open(img_byte_arr)
     except Exception as e:
         print(f"Watermark error: {str(e)}")
@@ -64,48 +61,43 @@ def add_watermark(image):
 # ===== IMAGE GENERATION =====
 def generate_image(prompt):
     if not prompt.strip():
-        return None, "⚠️ Please enter a prompt"
+        raise gr.Error("Please enter a prompt")
     
     try:
         model = load_model()
-        image = model(
+        result = model(
             prompt,
             num_inference_steps=30,
-            guidance_scale=7.5
-        ).images[0]
-        
-        return add_watermark(image), "✔️ Generation successful"
+            guidance_scale=7.5,
+            width=1024,
+            height=1024
+        )
+        return add_watermark(result.images[0]), "🎨 Generation complete!"
     
     except torch.cuda.OutOfMemoryError:
-        return None, "⚠️ Out of memory! Try a simpler prompt"
+        raise gr.Error("Out of memory! Try a simpler prompt or smaller image size")
     except Exception as e:
-        return None, f"⚠️ Error: {str(e)[:200]}"
+        raise gr.Error(f"Generation failed: {str(e)[:200]}")
 
-# ===== GRADIO THEME =====
-theme = gr.themes.Default(
+# ===== GRADIO 5.x INTERFACE =====
+with gr.Blocks(theme=gr.themes.Default(
     primary_hue="emerald",
-    secondary_hue="amber",
+    secondary_hue="gold",
     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
-    *Powered by HiDream-I1-Full (1024x1024 PNG output)*
-    """)
+)) as demo:
     
-    with gr.Row():
+    gr.Markdown("""<h1 align="center">🎨 SelamGPT HiDream Generator</h1>""")
+    
+    with gr.Row(variant="panel"):
         with gr.Column(scale=3):
             prompt_input = gr.Textbox(
                 label="Describe your image",
                 placeholder="A futuristic Ethiopian city with flying cars...",
                 lines=3,
-                max_lines=5
+                max_lines=5,
+                autofocus=True
             )
-            with gr.Row():
-                generate_btn = gr.Button("Generate Image", variant="primary")
-                clear_btn = gr.Button("Clear")
+            generate_btn = gr.Button("Generate Image", variant="primary")
             
             gr.Examples(
                 examples=[
@@ -113,33 +105,41 @@ with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
                     ["Traditional Ethiopian coffee ceremony in zero gravity"],
                     ["Portrait of a Habesha queen with golden jewelry"]
                 ],
-                inputs=prompt_input
+                inputs=prompt_input,
+                label="Try these prompts:"
             )
-            
+        
         with gr.Column(scale=2):
             output_image = gr.Image(
                 label="Generated Image",
                 type="pil",
-                format="png",
-                height=512
+                height=512,
+                interactive=False
             )
-            status_output = gr.Textbox(
+            status = gr.Textbox(
                 label="Status",
-                interactive=False
+                interactive=False,
+                show_label=False
             )
     
+    # Event handlers
     generate_btn.click(
         fn=generate_image,
         inputs=prompt_input,
-        outputs=[output_image, status_output],
-        queue=True
+        outputs=[output_image, status],
+        api_name="generate"
     )
     
-    clear_btn.click(
-        fn=lambda: [None, ""],
-        outputs=[output_image, status_output]
+    # Keyboard shortcut
+    prompt_input.submit(
+        fn=generate_image,
+        inputs=prompt_input,
+        outputs=[output_image, status]
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=2)
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )