Update app.py

app.py

@@ -12,22 +12,27 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 TORCH_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
 # ===== MODEL LOADING =====
-@gr.Cache()  # Cache model between generations
+# Global variable for model caching (alternative to @gr.Cache)
+pipe = None
+
 def load_model():
-    pipe = DiffusionPipeline.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=TORCH_DTYPE
-    ).to(DEVICE)
-    
-    # Optimizations
-    if DEVICE == "cuda":
-        pipe.enable_xformers_memory_efficient_attention()
-        pipe.enable_attention_slicing()
+    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()
     
     return pipe
 
-pipe = load_model()
-
 # ===== WATERMARK FUNCTION =====
 def add_watermark(image):
     """Add watermark with optimized PNG output"""
@@ -62,81 +67,19 @@ def generate_image(prompt):
         return None, "⚠️ Please enter a prompt"
     
     try:
-        # Generate image (1024x1024 by default)
-        image = pipe(
+        model = load_model()
+        image = model(
             prompt,
             num_inference_steps=30,
             guidance_scale=7.5
         ).images[0]
         
-        # Add watermark
-        watermarked = add_watermark(image)
-        return watermarked, "✔️ Generation successful"
+        return add_watermark(image), "✔️ Generation successful"
     
     except torch.cuda.OutOfMemoryError:
         return None, "⚠️ Out of memory! Try a simpler prompt"
     except Exception as e:
         return None, f"⚠️ Error: {str(e)[:200]}"
 
-# ===== 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
-    *Powered by HiDream-I1-Full (1024x1024 PNG output)*
-    """)
-    
-    with gr.Row():
-        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
-            )
-            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, 4k detailed"],
-                    ["Traditional Ethiopian coffee ceremony in zero gravity"],
-                    ["Portrait of a Habesha queen with golden jewelry"]
-                ],
-                inputs=prompt_input
-            )
-            
-        with gr.Column(scale=2):
-            output_image = gr.Image(
-                label="Generated Image",
-                type="pil",
-                format="png",
-                height=512
-            )
-            status_output = gr.Textbox(
-                label="Status",
-                interactive=False
-            )
-    
-    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__":
-    demo.queue(max_size=2)
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+# ... (keep your existing interface code exactly as is) ...