app.py

# -*- coding: utf-8 -*-
# ZenCtrl Inpainting Playground (Baseten backend)

#import spaces
import os, json, base64, requests
from io import BytesIO
from PIL import Image, ImageDraw
import gradio as gr
import replicate

# ────────── Secrets & endpoints ──────────
BASETEN_MODEL_URL = os.getenv("BASETEN_MODEL_URL")
BASETEN_API_KEY = os.getenv("BASETEN_API_KEY")
REPLICATE_TOKEN = os.getenv("REPLICATE_API_TOKEN")

# ────────── Globals ──────────
ADAPTER_SIZE = 1024
css = "#col-container {margin:0 auto; max-width:960px;}"

# Background generation via Replicate
def _gen_bg(prompt: str):
    url = replicate.run(
        "google/imagen-4-fast",
        input={"prompt": prompt or "cinematic background", "aspect_ratio": "1:1"},
    )
    url = url[0] if isinstance(url, list) else url
    return Image.open(BytesIO(requests.get(url, timeout=120).content)).convert("RGB")

# Main processing function
def process_image_and_text(subject_image, adapter_dict, prompt, _unused1, _unused2, size=ADAPTER_SIZE, rank=10.0):
    seed, guidance_scale, steps = 42, 2.5, 28

    adapter_image = adapter_dict["image"] if isinstance(adapter_dict, dict) else adapter_dict
    if isinstance(adapter_dict, dict) and adapter_dict.get("mask") is not None:
        m = adapter_dict["mask"].convert("L").point(lambda p: 255 if p else 0)
        if bbox := m.getbbox():
            rect = Image.new("L", m.size, 0)
            ImageDraw.Draw(rect).rectangle(bbox, fill=255)
            m = rect
        green = Image.new("RGB", adapter_image.size, "#00FF00")
        adapter_image = Image.composite(green, adapter_image, m)

    def prep(img: Image.Image):
        w, h = img.size
        m = min(w, h)
        return img.crop(((w - m) // 2, (h - m) // 2, (w + m) // 2, (h + m) // 2)).resize((size, size), Image.LANCZOS)

    subj_proc = prep(subject_image)
    adap_proc = prep(adapter_image)

    def b64(img):
        buf = BytesIO()
        img.save(buf, format="PNG")
        return base64.b64encode(buf.getvalue()).decode()

    payload = {
        "prompt": prompt,
        "subject_image": b64(subj_proc),
        "adapter_image": b64(adap_proc),
        "height": size,
        "width": size,
        "steps": steps,
        "seed": seed,
        "guidance_scale": guidance_scale,
        "rank": rank,
    }

    headers = {"Content-Type": "application/json"}
    if BASETEN_API_KEY:
        headers["Authorization"] = f"Api-Key {BASETEN_API_KEY}"

    resp = requests.post(BASETEN_MODEL_URL, headers=headers, json=payload, timeout=180)
    resp.raise_for_status()
    data = resp.json()
    
    # Extract base64 image from 'blended' key
    if "blended" in data:
        try:
            blended_bytes = base64.b64decode(data["blended"])
            raw_img = Image.open(BytesIO(blended_bytes)).convert("RGB")
        except Exception:
            raise gr.Error("Failed to decode 'blended' image from Baseten response.")
    else:
        raise gr.Error("Baseten response missing 'blended' image.")



# ────────── Header HTML ──────────
header_html = """
<h1>ZenCtrl Inpainting</h1>
<div align="center" style="line-height:1;">
  <a href="https://discord.com/invite/b9RuYQ3F8k" target="_blank" style="margin:10px;">
    <img src="https://img.shields.io/badge/Discord-Join-7289da.svg?logo=discord" alt="Discord">
  </a>
  <a href="https://fotographer.ai/zen-control" target="_blank" style="margin:10px;">
    <img src="https://img.shields.io/badge/Website-Landing_Page-blue" alt="LP">
  </a>
  <a href="https://x.com/FotographerAI" target="_blank" style="margin:10px;">
    <img src="https://img.shields.io/twitter/follow/FotographerAI?style=social" alt="X">
  </a>
</div>
"""

# ────────── Gradio UI ──────────
with gr.Blocks(css=css, title="ZenCtrl Inpainting") as demo:
    raw_state = gr.State()

    gr.HTML(header_html)
    gr.Markdown(
        "**Generate context-aware images of your subject with ZenCtrl’s inpainting playground.**  \n"
        "Upload a subject + optional mask, write a prompt, and hit **Generate**.  \n"
        "Open *Advanced Settings* to fetch an AI-generated background."
    )

    with gr.Row():
        with gr.Column(scale=2, elem_id="col-container"):
            subj_img = gr.Image(type="pil", label="Subject image")
            ref_img = gr.Image(type="pil", label="Background / Mask image", tool="sketch", brush_color="#00FF00")
            ref_img_ex = gr.Image(type="pil", visible=False)

            # Removed Florence-SAM
            promptbox = gr.Textbox(label="Generation prompt", value="furniture", lines=2)
            run_btn = gr.Button("Generate", variant="primary")

            with gr.Accordion("Advanced Settings", open=False):
                bgprompt = gr.Textbox(label="Background Prompt", value="Scandinavian living room …")
                bg_btn = gr.Button("Generate BG")

        with gr.Column(scale=2):
            gallery = gr.Gallery(columns=[1], rows=[1], object_fit="contain", height="auto")
            bg_img = gr.Image(label="Background", visible=True)

    # ---------- Example wrapper ---------------------------------
    def _run_example(subj, bg, prompt):
        if isinstance(subj, str):
            subj = Image.open(subj)
        if isinstance(bg, str):
            bg = Image.open(bg)

        adapter_dict = {"image": bg, "mask": None}
        gallery_out, _ = process_image_and_text(subj, adapter_dict, prompt, False, "")
        return gallery_out, gr.update(value=bg)

    # ---------- Examples ----------------------------------------
    gr.Examples(
        examples=[
            ["examples/subject1.png", "examples/subject1.png", "Make the toy sit on a marble table"],
            ["examples/subject1.png", "examples/subject1.png", "Turn the flowers into sunflowers"],
            ["examples/subject1.png", "examples/subject1.png", "Make this monster ride a skateboard on the beach"],
            ["examples/subject1.png", "examples/subject1.png", "Make this cat happy"],
        ],
        inputs=[subj_img, ref_img_ex, promptbox],
        outputs=[gallery, ref_img],
        fn=_run_example,
        examples_per_page="all",
        label="Presets (Input · Background · Prompt)",
        cache_examples=False,
    )

    # ---------- Buttons & interactions --------------------------
    run_btn.click(
        process_image_and_text,
        inputs=[subj_img, ref_img, promptbox, gr.State(False), gr.State("")],
        outputs=[gallery, raw_state]
    )
    bg_btn.click(_gen_bg, inputs=[bgprompt], outputs=[bg_img])

# ---------------- Launch ---------------------------------------
if __name__ == "__main__":
    demo.launch(debug=True, share=True)