Update app.py

app.py

@@ -4,44 +4,126 @@ import requests
 import gradio as gr
 from PIL import Image
 from io import BytesIO
+from gradio_client import Client
 
-def generate_image(prompt, width, height, seed, randomize):
-    print(prompt+"\n\n")
+# more servers coming soon...
+
+
+SERVER_NAMES = {
+    "google_us": "Google US Server",
+    "azure_lite": "Azure Lite Supercomputer Server",
+    
+}
+
+
+SERVER_SOCKETS = {
+    "google_us": None,
+    "azure_lite": "FLUX-Pro-SERVER1",
+    
+}
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+FLUX_URL  = os.environ.get("FLUX_URL")
+
+
+
+def _open_image_from_str(s: str):
+    # base64 decoding
+    if s.startswith("http"):
+        r = requests.get(s); return Image.open(BytesIO(r.content))
+    if os.path.exists(s):
+        return Image.open(s)
+    # try base64 blob
+    try:
+        import base64
+        _, b64 = s.split(",", 1)
+        data = base64.b64decode(b64)
+        return Image.open(BytesIO(data))
+    except:
+        raise ValueError(f"Can't parse image string: {s[:30]}…")
+
+
+
+def generate_image(prompt, width, height, seed, randomize, server_choice):
+
+    print("\n\n\n\n"+prompt+"\n\n\n\n")
+    
     if randomize:
-        seed = random.randint(0, 9999999)
-    base_url = os.environ.get("FLUX_URL")
-    base_url = base_url + ""
-    if not base_url:
-        return "Error: FLUX_URL environment variable not set."
-    url = (
-        base_url.replace("[prompt]", prompt)
-        .replace("[w]", str(width))
-        .replace("[h]", str(height))
-        .replace("[seed]", str(seed))
+        seed = random.randint(0, 9_999_999)
+
+    # dict key for this display name
+    key = next(k for k, v in SERVER_NAMES.items() if v == server_choice)
+    socket = SERVER_SOCKETS.get(key)
+
+    
+    if socket is None:
+        if not FLUX_URL:
+            return "Error: FLUX_URL not set."
+        url = (
+            FLUX_URL
+            .replace("[prompt]", prompt)
+            .replace("[w]", str(width))
+            .replace("[h]", str(height))
+            .replace("[seed]", str(seed))
+        )
+        r = requests.get(url)
+        return Image.open(BytesIO(r.content)) if r.ok else f"FLUX-Pro failed ({r.status_code})"
+
+    
+    space_id = f"NihalGazi/{socket}"
+    client = Client(space_id, hf_token=HF_TOKEN)
+    res = client.predict(
+        prompt=prompt,
+        width=width,
+        height=height,
+        seed=seed,
+        randomize=randomize,
+        api_name="/predict"
     )
-    response = requests.get(url)
-    if response.status_code == 200:
-        image = Image.open(BytesIO(response.content))
-        return image
+
+    # handle dict or str
+    if isinstance(res, dict):
+        if res.get("path"):
+            return Image.open(res["path"])
+        if res.get("url"):
+            return _open_image_from_str(res["url"])
+        return "No image found in response."
+    elif isinstance(res, str):
+        return _open_image_from_str(res)
     else:
-        return "Failed to retrieve image."
+        return f"Unexpected response type: {type(res)}"
 
-title = "Unlimited FLUX-Pro"
-description = "Here you can use the FLUX-Pro as much as you want without limits"
+# ─── GRADIO INTERFACE ─────────────────────────────────────────────────────
 
 iface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your image prompt..."),
-        gr.Slider(minimum=512, maximum=1280, step=16, value=1280, label="Width"),
-        gr.Slider(minimum=512, maximum=1280, step=16, value=1280, label="Height"),
+        gr.Textbox(label="Prompt", placeholder="Enter your image prompt…"),
+        gr.Slider(512, 1280, step=16, value=1280, label="Width"),
+        gr.Slider(512, 1280, step=16, value=1280, label="Height"),
         gr.Number(label="Seed", value=0),
-        gr.Checkbox(label="Randomize Seed", value=True)
+        gr.Checkbox(label="Randomize Seed", value=True),
+        gr.Dropdown(
+            label="Server",
+            choices=list(SERVER_NAMES.values()),
+            value=SERVER_NAMES["google_us"]
+        ),
     ],
     outputs=gr.Image(type="pil"),
-    title=title,
-    description=description,
+    title="Unlimited FLUX‑Pro",
+    description="""**Enter a prompt and tweak your settings:**  
+-  **Server** – switch between servers *if one is slow or fails*:
+    1. **Google US Server**
+    2. **Azure Lite Supercomputer Server**
+- **Width & Height** – choose your canvas size  
+- **Seed** – pick a number or check **Randomize Seed**  
+
+Click **Generate** and enjoy unlimited AI art!
+❤️ **Like & follow** for more AI projects:  
+• Instagram: [@nihal_gazi_io](https://www.instagram.com/nihal_gazi_io/)  
+• Discord: nihal_gazi_io
+    """
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()