app.py

import os
import random
import requests
import gradio as gr
from PIL import Image
from io import BytesIO
from gradio_client import Client

# more servers coming soon...


SERVER_NAMES = {
    "google_us": "Google US Server",
    "azure_lite": "Azure Lite Supercomputer Server",
    "artemis" : "Artemis GPU Super cluster",
    "nsfw_core" : "NSFW-Core: Uncensored Server",
    "nsfw_core_2" : "NSFW-Core: Uncensored Server 2",
    
}


SERVER_SOCKETS = {
    "google_us": None,
    "azure_lite": "FLUX-Pro-SERVER1",
    "artemis" : "FLUX-Pro-Artemis-GPU",
    "nsfw_core": "FLUX-Pro-NSFW-LocalCoreProcessor",
    "nsfw_core_2" : "FLUX-Pro-NSFW-LocalCoreProcessor-v2",
    
}

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, 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"
    )

    # 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 f"Unexpected response type: {type(res)}"

# ─── GRADIO INTERFACE ─────────────────────────────────────────────────────

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        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.Dropdown(
            label="Server",
            choices=list(SERVER_NAMES.values()),
            value=SERVER_NAMES["google_us"]
        ),
    ],
    outputs=gr.Image(type="pil"),
    title="Unlimited FLUX‑Pro",
    description = """
**Enter a prompt and tweak your settings:**  
- **Width & Height** – choose your canvas size  
- **Seed** – pick a number or check **Randomize Seed**  
- **Server** – switch between servers if one is slow or fails:
  - **Google US Server**  
  - **Azure Lite Supercomputer Server**  
  - **Artemis GPU Super cluster**  
  - **NSFW‑Core: Uncensored Servers** (for explicit content; use responsibly)
- **Suggestions** – have ideas? I’m open to them!

⚠️ **Caution:**  
The **NSFW‑Core** server can generate adult‑only content. You must be of legal age in your jurisdiction and comply with all local laws and platform policies. Developer is not liable for misuse.


> ⚡ 2 NSFW Servers available 


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()