# -*- 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 = """

ZenCtrl Inpainting

Discord LP X
""" # ────────── 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)