Upload app.py with huggingface_hub

app.py

@@ -73,46 +73,57 @@ def query_databank(filename: str) -> str:
 
 @mcp.tool()
 def generate_image(prompt: str) -> str:
-    """Generate an image with multiple model fallbacks using the free hf-inference provider."""
+    """Generate an image using a free ZeroGPU Space via gradio_client, with fallbacks."""
     try:
-        from huggingface_hub import InferenceClient
-        hf_client = InferenceClient(
-            provider="hf-inference", 
-            token=HF_TOKEN,
-            headers={"x-wait-for-model": "true"}
-        )
+        from gradio_client import Client
+        import shutil
         
         business_name = os.environ.get("BUSINESS_NAME", "Fair Dinkum Publishing")
         owner = os.environ.get("BUSINESS_OWNER", "BRETT SJOBERG")
         brand_context = f"Professional brand asset for {business_name} (Owner: {owner}). Style: Modern, clean, high-quality. "
         full_prompt = brand_context + prompt
         
-        # Truly free models on hf-inference
-        models = [
-            "runwayml/stable-diffusion-v1-5",
-            "stabilityai/stable-diffusion-2-1",
-            "black-forest-labs/FLUX.1-schnell"
-        ]
-        
-        image = None
-        model_used = None
-        
-        for model_id in models:
-            try:
-                image = hf_client.text_to_image(full_prompt, model=model_id)
-                model_used = model_id
-                break
-            except Exception as e:
-                print(f"Model {model_id} failed: {str(e)}")
-                continue
-        
-        if not image:
-            return "Error: All free image generation models failed. The serverless API may be overloaded."
-                
-        os.makedirs("exports/images", exist_ok=True)
-        image_path = f"exports/images/{abs(hash(prompt))}.png"
-        image.save(image_path)
-        return f"Branded Image Generated using {model_used} (Free Tier): {image_path}"
+        try:
+            # Method 1: Gradio Client (ZeroGPU Space - Truly Free)
+            print(f"Attempting free generation via Gradio Client...")
+            client = Client("black-forest-labs/FLUX.1-schnell", hf_token=HF_TOKEN)
+            result = client.predict(
+                prompt=full_prompt,
+                seed=0,
+                randomize_seed=True,
+                width=1024,
+                height=1024,
+                num_inference_steps=4,
+                api_name="/infer"
+            )
+            # Result is a tuple where the first element is the image path
+            temp_image_path = result[0] if isinstance(result, (list, tuple)) else result
+            
+            os.makedirs("exports/images", exist_ok=True)
+            final_path = f"exports/images/{abs(hash(prompt))}.png"
+            shutil.copy(temp_image_path, final_path)
+            return f"Branded Image Generated using FLUX.1 (Free ZeroGPU Space): {final_path}"
+            
+        except Exception as g_e:
+            print(f"Gradio Client method failed: {str(g_e)}. Falling back to classic serverless...")
+            
+            # Fallback to Classic Serverless (Method 2)
+            from huggingface_hub import InferenceClient
+            hf_client = InferenceClient(provider="hf-inference", token=HF_TOKEN, headers={"x-wait-for-model": "true"})
+            
+            models = ["runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1"]
+            for model_id in models:
+                try:
+                    image = hf_client.text_to_image(full_prompt, model=model_id)
+                    os.makedirs("exports/images", exist_ok=True)
+                    image_path = f"exports/images/{abs(hash(prompt))}.png"
+                    image.save(image_path)
+                    return f"Branded Image Generated using {model_id} (Free Serverless Tier): {image_path}"
+                except Exception as e:
+                    continue
+            
+            return f"Error: All free generation methods failed. (Gradio Error: {str(g_e)})"
+            
     except Exception as e:
         return f"System Error during image generation: {str(e)}"