app.py

import os
import base64
import io
from typing import TypedDict
import requests
import gradio as gr
from PIL import Image

# Read Baseten configuration from environment variables.
BTEN_API_KEY = os.getenv("API_KEY")
URL = os.getenv("URL")


def image_to_base64(image: Image.Image) -> str:
    with io.BytesIO() as buffer:
        image.save(buffer, format="PNG")
        return base64.b64encode(buffer.getvalue()).decode("utf-8")


def ensure_image(img) -> Image.Image:
    if isinstance(img, Image.Image):
        return img
    elif isinstance(img, str):
        return Image.open(img)
    elif isinstance(img, dict) and "name" in img:
        return Image.open(img["name"])
    else:
        raise ValueError("Cannot convert input to a PIL Image.")


def call_baseten_generate(
    image: Image.Image,
    prompt: str,
    steps: int,
    strength: float,
    height: int,
    width: int,
    lora_name: str,
    remove_bg: bool,
) -> Image.Image | None:
    image = ensure_image(image)
    b64_image = image_to_base64(image)
    payload = {
        "image": b64_image,
        "prompt": prompt,
        "steps": steps,
        "strength": strength,
        "height": height,
        "width": width,
        "lora_name": lora_name,
        "bgrm": remove_bg,
    }
    headers = {"Authorization": f"Api-Key {BTEN_API_KEY or os.getenv('API_KEY')}"}
    try:
        if not URL:
            raise ValueError("The URL environment variable is not set.")
        response = requests.post(URL, headers=headers, json=payload)
        if response.status_code == 200:
            data = response.json()
            gen_b64 = data.get("generated_image", None)
            if gen_b64:
                return Image.open(io.BytesIO(base64.b64decode(gen_b64)))
            else:
                return None
        else:
            print(f"Error: HTTP {response.status_code}\n{response.text}")
            return None
    except Exception as e:
        print(f"Error: {e}")
        return None


# ================== MODE CONFIG =====================

Mode = TypedDict(
    "Mode",
    {
        "model": str,
        "prompt": str,
        "default_strength": float,
        "default_height": int,
        "default_width": int,
        "models": list[str],
        "remove_bg": bool,
    },
)

MODE_DEFAULTS: dict[str, Mode] = {
    "Subject Generation": {
        "model": "subject_99000_512",
        "prompt": "A detailed portrait with soft lighting",
        "default_strength": 1.2,
        "default_height": 512,
        "default_width": 512,
        "models": ["zendsd_512_146000", "subject_99000_512", "zen_26000_512"],
        "remove_bg": True,
    },
    "Background Generation": {
        "model": "bg_canny_58000_1024",
        "prompt": "A vibrant background with dynamic lighting and textures",
        "default_strength": 1.2,
        "default_height": 1024,
        "default_width": 1024,
        "models": ["bgwlight_15000_1024", "bg_canny_58000_1024", "gen_back_7000_1024"],
        "remove_bg": True,
    },
    "Canny": {
        "model": "canny_21000_1024",
        "prompt": "A futuristic cityscape with neon lights",
        "default_strength": 1.2,
        "default_height": 1024,
        "default_width": 1024,
        "models": ["canny_21000_1024"],
        "remove_bg": True,
    },
    "Depth": {
        "model": "depth_9800_1024",
        "prompt": "A scene with pronounced depth and perspective",
        "default_strength": 1.2,
        "default_height": 1024,
        "default_width": 1024,
        "models": ["depth_9800_1024"],
        "remove_bg": True,
    },
    "Deblurring": {
        "model": "deblurr_1024_10000",
        "prompt": "A scene with pronounced depth and perspective",
        "default_strength": 1.2,
        "default_height": 1024,
        "default_width": 1024,
        "models": ["deblurr_1024_10000"],
        "remove_bg": False,
    },
}

# ================== PRESET EXAMPLES =====================

MODE_EXAMPLES = {
    "Subject Generation": [
        [
            "assets/subj1.jpg",
            "Close-up portrait of a fruit bowl",
            "assets/subj1_out.jpg",
        ],
        ["assets/subj2.jpg", "A penguin standing in snow", "assets/subj2_out.jpg"],
        ["assets/subj3.jpg", "A cat with glowing eyes", "assets/subj3_out.jpg"],
        ["assets/subj4.jpg", "A child playing with bubbles", "assets/subj4_out.jpg"],
        [
            "assets/subj5.jpg",
            "A stylish young man in neon lights",
            "assets/subj5_out.jpg",
        ],
        ["assets/subj6.jpg", "Old man with a mysterious look", "assets/subj6_out.jpg"],
    ],
    "Background Generation": [
        ["assets/bg1.jpg", "Modern living room with plants", "assets/bg1_out.jpg"],
        ["assets/bg2.jpg", "Fantasy forest background", "assets/bg2_out.jpg"],
        ["assets/bg3.jpg", "Futuristic cityscape", "assets/bg3_out.jpg"],
        ["assets/bg4.jpg", "Minimalist white studio", "assets/bg4_out.jpg"],
        ["assets/bg5.jpg", "Snowy mountain landscape", "assets/bg5_out.jpg"],
        ["assets/bg6.jpg", "Golden sunset over the sea", "assets/bg6_out.jpg"],
    ],
    "Canny": [
        ["assets/canny1.jpg", "A neon cyberpunk city skyline", "assets/canny1_out.jpg"],
        ["assets/canny2.jpg", "A robot walking in the fog", "assets/canny2_out.jpg"],
        [
            "assets/canny3.jpg",
            "A futuristic vehicle parked under a bridge",
            "assets/canny3_out.jpg",
        ],
        [
            "assets/canny4.jpg",
            "Sci-fi lab interior with glowing machinery",
            "assets/canny4_out.jpg",
        ],
        [
            "assets/canny5.jpg",
            "A portrait of a woman outlined in neon",
            "assets/canny5_out.jpg",
        ],
        [
            "assets/canny6.jpg",
            "Post-apocalyptic abandoned street",
            "assets/canny6_out.jpg",
        ],
    ],
    "Depth": [
        [
            "assets/depth1.jpg",
            "A narrow alleyway with deep perspective",
            "assets/depth1_out.jpg",
        ],
        [
            "assets/depth2.jpg",
            "A mountain road vanishing into the distance",
            "assets/depth2_out.jpg",
        ],
        [
            "assets/depth3.jpg",
            "A hallway with strong depth of field",
            "assets/depth3_out.jpg",
        ],
        [
            "assets/depth4.jpg",
            "A misty forest path stretching far away",
            "assets/depth4_out.jpg",
        ],
        ["assets/depth5.jpg", "A bridge over a deep canyon", "assets/depth5_out.jpg"],
        [
            "assets/depth6.jpg",
            "An underground tunnel with receding arches",
            "assets/depth6_out.jpg",
        ],
    ],
    "Deblurring": [
        ["assets/deblur1.jpg", "", "assets/deblur1_out.jpg"],
        ["assets/deblur2.jpg", "", "assets/deblur2_out.jpg"],
        ["assets/deblur3.jpg", "", "assets/deblur3_out.jpg"],
        ["assets/deblur4.jpg", "", "assets/deblur4_out.jpg"],
        ["assets/deblur5.jpg", "", "assets/deblur5_out.jpg"],
        ["assets/deblur6.jpg", "", "assets/deblur6_out.jpg"],
    ],
}


# ================== UI =====================

header = """
<h1>🌍 ZenCtrl / FLUX</h1>
<div align="center" style="line-height: 1;">
    <a href="https://github.com/FotographerAI/ZenCtrl/tree/main" target="_blank"><img src="https://img.shields.io/badge/GitHub-Repo-181717.svg"></a>
    <a href="https://huggingface.co/spaces/fotographerai/ZenCtrl" target="_blank"><img src="https://img.shields.io/badge/🤗_HuggingFace-Space-ffbd45.svg"></a>
    <a href="https://discord.com/invite/b9RuYQ3F8k" target="_blank"><img src="https://img.shields.io/badge/Discord-Join-7289da.svg?logo=discord"></a>
</div>
"""

with gr.Blocks(title="🌍 ZenCtrl") as demo:
    gr.HTML(header)
    gr.Markdown("# ZenCtrl Demo")

    with gr.Tabs():
        for mode in MODE_DEFAULTS:
            with gr.Tab(mode):
                defaults = MODE_DEFAULTS[mode]
                gr.Markdown(f"### {mode} Mode")

                with gr.Row():
                    with gr.Column(scale=2):
                        input_image = gr.Image(label="Input Image", type="pil")
                        generate_button = gr.Button("Generate")
                        with gr.Blocks():
                            model_dropdown = gr.Dropdown(
                                label="Model",
                                choices=defaults["models"],
                                value=defaults["model"],
                                interactive=True,
                            )
                            remove_bg_checkbox = gr.Checkbox(
                                label="Remove Background", value=defaults["remove_bg"]
                            )

                    with gr.Column(scale=2):
                        output_image = gr.Image(label="Generated Image", type="pil")

                prompt_box = gr.Textbox(
                    label="Prompt", value=defaults["prompt"], lines=2
                )

                with gr.Accordion("Generation Parameters", open=False):
                    with gr.Row():
                        step_slider = gr.Slider(
                            minimum=2, maximum=28, value=2, step=2, label="Steps"
                        )
                        strength_slider = gr.Slider(
                            minimum=0.5,
                            maximum=2.0,
                            value=defaults["default_strength"],
                            step=0.1,
                            label="Strength",
                        )
                    with gr.Row():
                        height_slider = gr.Slider(
                            minimum=512,
                            maximum=1360,
                            value=defaults["default_height"],
                            step=1,
                            label="Height",
                        )
                        width_slider = gr.Slider(
                            minimum=512,
                            maximum=1360,
                            value=defaults["default_width"],
                            step=1,
                            label="Width",
                        )

                def on_generate_click(
                    model_name, prompt, steps, strength, height, width, remove_bg, image
                ):
                    return call_baseten_generate(
                        image,
                        prompt,
                        steps,
                        strength,
                        height,
                        width,
                        model_name,
                        remove_bg,
                    )

                generate_button.click(
                    fn=on_generate_click,
                    inputs=[
                        model_dropdown,
                        prompt_box,
                        step_slider,
                        strength_slider,
                        height_slider,
                        width_slider,
                        remove_bg_checkbox,
                        input_image,
                    ],
                    outputs=[output_image],
                )

                # ---------------- Templates --------------------
                gr.Dataset(
                    label="Presets (Input / Prompt / Output)",
                    headers=["Input", "Prompt", "Output"],
                    components=[input_image, prompt_box, output_image],
                    samples=MODE_EXAMPLES.get(mode, []),
                    samples_per_page=6,
                )

if __name__ == "__main__":
    demo.launch()