app.py

import json
import os
import random
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path

import gradio as gr
import requests
import spaces
from huggingface_hub import hf_hub_download
from PIL import Image


MODEL_REPO = "Daankular/redcraft-krea2-fp8"
MODEL_FILE = "redcraftKREA2RedMix_krea2Edition.safetensors"
COMFY_KREA_REPO = "Comfy-Org/Krea-2"
TEXT_ENCODER_FILE = "qwen3vl_4b_bf16.safetensors"
VAE_FILE = "qwen_image_vae.safetensors"
IDENTITY_LORA_REPO = "conradlocke/krea2-identity-edit"
IDENTITY_LORA_FILE = "krea2_identity_edit_v1_2.safetensors"
KREA2EDIT_NODE_REPO = "https://github.com/lbouaraba/comfyui-krea2edit"
KREA2EDIT_NODE_DIRNAME = "comfyui-krea2edit"
COMFY_REPO = "https://github.com/comfyanonymous/ComfyUI.git"
COMFY_DIR = Path(os.environ.get("COMFYUI_DIR", "/tmp/ComfyUI"))
COMFY_HOST = "127.0.0.1"
COMFY_PORT = int(os.environ.get("COMFYUI_PORT", "8188"))
COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
MAX_SEED = 2**31 - 1

_comfy_lock = threading.Lock()
_comfy_process = None


def _run(cmd, cwd=None):
    print("[setup]", " ".join(map(str, cmd)), flush=True)
    subprocess.check_call(cmd, cwd=str(cwd) if cwd else None)


def _wait_for_comfy(timeout=180):
    deadline = time.time() + timeout
    last_error = None
    while time.time() < deadline:
        try:
            response = requests.get(f"{COMFY_URL}/system_stats", timeout=2)
            if response.ok:
                return
        except Exception as exc:
            last_error = exc
        time.sleep(1)
    raise RuntimeError(f"ComfyUI did not start in time: {last_error}")


def _validate_comfyui():
    response = requests.get(f"{COMFY_URL}/object_info", timeout=30)
    response.raise_for_status()
    object_info = response.json()
    required_nodes = ["UNETLoader", "CLIPLoader", "VAELoader", "CLIPTextEncode", "KSampler", "VAEDecode", "SaveImage"]
    missing = [node for node in required_nodes if node not in object_info]
    if missing:
        raise RuntimeError(f"ComfyUI is missing required nodes: {', '.join(missing)}")

    unet_info = object_info["UNETLoader"]["input"]["required"]["unet_name"][0]
    clip_info = object_info["CLIPLoader"]["input"]["required"]["clip_name"][0]
    vae_info = object_info["VAELoader"]["input"]["required"]["vae_name"][0]
    if MODEL_FILE not in unet_info:
        raise RuntimeError(f"Redcraft diffusion model is not visible to ComfyUI. First models: {', '.join(unet_info[:10])}")
    if TEXT_ENCODER_FILE not in clip_info:
        raise RuntimeError(f"Krea2 text encoder is not visible to ComfyUI. First encoders: {', '.join(clip_info[:10])}")
    if VAE_FILE not in vae_info:
        raise RuntimeError(f"Krea2 VAE is not visible to ComfyUI. First VAEs: {', '.join(vae_info[:10])}")

    identity_edit_nodes = ["LoraLoaderModelOnly", "Krea2EditModelPatch", "Krea2EditGroundedEncode", "EmptySD3LatentImage"]
    missing_identity_nodes = [node for node in identity_edit_nodes if node not in object_info]
    if missing_identity_nodes:
        raise RuntimeError(
            f"ComfyUI-Krea2Edit nodes are missing: {', '.join(missing_identity_nodes)}. "
            f"Check the {KREA2EDIT_NODE_REPO} custom node install."
        )
    lora_info = object_info["LoraLoaderModelOnly"]["input"]["required"]["lora_name"][0]
    if IDENTITY_LORA_FILE not in lora_info:
        raise RuntimeError(f"Identity-edit LoRA is not visible to ComfyUI. First loras: {', '.join(lora_info[:10])}")


def _ensure_comfyui():
    global _comfy_process

    with _comfy_lock:
        if _comfy_process is not None and _comfy_process.poll() is None:
            return
        try:
            response = requests.get(f"{COMFY_URL}/system_stats", timeout=2)
            if response.ok:
                return
        except Exception:
            pass

        if not COMFY_DIR.exists():
            _run(["git", "clone", "--depth", "1", COMFY_REPO, str(COMFY_DIR)])

        marker = COMFY_DIR / ".requirements-installed"
        if not marker.exists():
            _run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=COMFY_DIR)
            marker.write_text("ok", encoding="utf-8")

        krea2edit_dir = COMFY_DIR / "custom_nodes" / KREA2EDIT_NODE_DIRNAME
        if not krea2edit_dir.exists():
            _run(["git", "clone", "--depth", "1", KREA2EDIT_NODE_REPO, str(krea2edit_dir)])

        diffusion_dir = COMFY_DIR / "models" / "diffusion_models"
        text_encoder_dir = COMFY_DIR / "models" / "text_encoders"
        vae_dir = COMFY_DIR / "models" / "vae"
        lora_dir = COMFY_DIR / "models" / "loras"
        diffusion_dir.mkdir(parents=True, exist_ok=True)
        text_encoder_dir.mkdir(parents=True, exist_ok=True)
        vae_dir.mkdir(parents=True, exist_ok=True)
        lora_dir.mkdir(parents=True, exist_ok=True)
        hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            local_dir=str(diffusion_dir),
            token=os.environ.get("HF_TOKEN"),
        )
        hf_hub_download(
            repo_id=COMFY_KREA_REPO,
            filename=f"text_encoders/{TEXT_ENCODER_FILE}",
            local_dir=str(COMFY_DIR / "models"),
            token=os.environ.get("HF_TOKEN"),
        )
        hf_hub_download(
            repo_id=COMFY_KREA_REPO,
            filename=f"vae/{VAE_FILE}",
            local_dir=str(COMFY_DIR / "models"),
            token=os.environ.get("HF_TOKEN"),
        )
        hf_hub_download(
            repo_id=IDENTITY_LORA_REPO,
            filename=IDENTITY_LORA_FILE,
            local_dir=str(lora_dir),
            token=os.environ.get("HF_TOKEN"),
        )


def _start_comfyui():
    global _comfy_process

    with _comfy_lock:
        if _comfy_process is not None and _comfy_process.poll() is None:
            return
        try:
            response = requests.get(f"{COMFY_URL}/system_stats", timeout=2)
            if response.ok:
                _validate_comfyui()
                return
        except Exception:
            pass

        cmd = [
            sys.executable,
            "main.py",
            "--listen",
            COMFY_HOST,
            "--port",
            str(COMFY_PORT),
            "--disable-auto-launch",
        ]
        _comfy_process = subprocess.Popen(cmd, cwd=str(COMFY_DIR))
        _wait_for_comfy()
        _validate_comfyui()


def _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler):
    workflow = {
        "1": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": MODEL_FILE, "weight_dtype": "default"},
        },
        "8": {
            "class_type": "CLIPLoader",
            "inputs": {"clip_name": TEXT_ENCODER_FILE, "type": "krea2", "device": "default"},
        },
        "9": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": VAE_FILE},
        },
        "2": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": prompt, "clip": ["8", 0]},
        },
        "4": {
            "class_type": "EmptyLatentImage",
            "inputs": {"width": int(width), "height": int(height), "batch_size": 1},
        },
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["1", 0],
                "positive": ["2", 0],
                "negative": ["3", 0],
                "latent_image": ["4", 0],
            },
        },
        "6": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["5", 0], "vae": ["9", 0]},
        },
        "7": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "redcraft", "images": ["6", 0]},
        },
    }

    if negative_prompt and negative_prompt.strip():
        workflow["3"] = {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative_prompt, "clip": ["8", 0]},
        }
    else:
        workflow["3"] = {
            "class_type": "ConditioningZeroOut",
            "inputs": {"conditioning": ["2", 0]},
        }
    return workflow


def _upload_image_to_comfy(image):
    if image is None:
        raise ValueError("Upload an image to edit.")

    image = image.convert("RGB")
    filename = f"redcraft-input-{uuid.uuid4().hex}.png"
    temp_path = Path("/tmp") / filename
    image.save(temp_path)
    with temp_path.open("rb") as handle:
        response = requests.post(
            f"{COMFY_URL}/upload/image",
            files={"image": (filename, handle, "image/png")},
            data={"overwrite": "true"},
            timeout=120,
        )
    response.raise_for_status()
    data = response.json()
    return data.get("name", filename)


def _resize_for_edit(image, width, height):
    width, height = int(width), int(height)
    if width <= 0 or height <= 0:
        return image.convert("RGB")
    return image.convert("RGB").resize((width, height), Image.LANCZOS)


def _target_size_from_source(image, max_megapixels=1.0):
    multiple = 16
    width, height = image.size
    megapixels = (width * height) / 1_000_000
    if megapixels > max_megapixels:
        scale = (max_megapixels / megapixels) ** 0.5
        width, height = round(width * scale), round(height * scale)
    width = max(multiple, (width // multiple) * multiple)
    height = max(multiple, (height // multiple) * multiple)
    return width, height


def _build_edit_workflow(input_filename, prompt, negative_prompt, steps, cfg, seed, sampler, scheduler, denoise):
    workflow = {
        "1": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": MODEL_FILE, "weight_dtype": "default"},
        },
        "8": {
            "class_type": "CLIPLoader",
            "inputs": {"clip_name": TEXT_ENCODER_FILE, "type": "krea2", "device": "default"},
        },
        "9": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": VAE_FILE},
        },
        "10": {
            "class_type": "LoadImage",
            "inputs": {"image": input_filename},
        },
        "11": {
            "class_type": "VAEEncode",
            "inputs": {"pixels": ["10", 0], "vae": ["9", 0]},
        },
        "2": {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": prompt, "clip": ["8", 0]},
        },
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler,
                "scheduler": scheduler,
                "denoise": float(denoise),
                "model": ["1", 0],
                "positive": ["2", 0],
                "negative": ["3", 0],
                "latent_image": ["11", 0],
            },
        },
        "6": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["5", 0], "vae": ["9", 0]},
        },
        "7": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "redcraft-edit", "images": ["6", 0]},
        },
    }

    if negative_prompt and negative_prompt.strip():
        workflow["3"] = {
            "class_type": "CLIPTextEncode",
            "inputs": {"text": negative_prompt, "clip": ["8", 0]},
        }
    else:
        workflow["3"] = {
            "class_type": "ConditioningZeroOut",
            "inputs": {"conditioning": ["2", 0]},
        }
    return workflow


def _build_identity_edit_workflow(input_filename, prompt, width, height, steps, cfg, seed, ref_boost, grounding_px, sampler, scheduler):
    return {
        "1": {
            "class_type": "UNETLoader",
            "inputs": {"unet_name": MODEL_FILE, "weight_dtype": "default"},
        },
        "8": {
            "class_type": "CLIPLoader",
            "inputs": {"clip_name": TEXT_ENCODER_FILE, "type": "krea2", "device": "default"},
        },
        "9": {
            "class_type": "VAELoader",
            "inputs": {"vae_name": VAE_FILE},
        },
        "20": {
            "class_type": "LoraLoaderModelOnly",
            "inputs": {"lora_name": IDENTITY_LORA_FILE, "strength_model": 1.0, "model": ["1", 0]},
        },
        "10": {
            "class_type": "LoadImage",
            "inputs": {"image": input_filename},
        },
        "11": {
            "class_type": "VAEEncode",
            "inputs": {"pixels": ["10", 0], "vae": ["9", 0]},
        },
        "21": {
            "class_type": "Krea2EditModelPatch",
            "inputs": {
                "model": ["20", 0],
                "source_latent": ["11", 0],
                "ref_boost": float(ref_boost),
                "fit_mode": "fit",
                "vae": ["9", 0],
                "source_image": ["10", 0],
            },
        },
        "2": {
            "class_type": "Krea2EditGroundedEncode",
            "inputs": {"clip": ["8", 0], "prompt": prompt, "image": ["10", 0], "grounding_px": int(grounding_px)},
        },
        "3": {
            "class_type": "Krea2EditGroundedEncode",
            "inputs": {"clip": ["8", 0], "prompt": "", "image": ["10", 0], "grounding_px": int(grounding_px)},
        },
        "4": {
            "class_type": "EmptySD3LatentImage",
            "inputs": {"width": int(width), "height": int(height), "batch_size": 1},
        },
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": int(seed),
                "steps": int(steps),
                "cfg": float(cfg),
                "sampler_name": sampler,
                "scheduler": scheduler,
                "denoise": 1.0,
                "model": ["21", 0],
                "positive": ["2", 0],
                "negative": ["3", 0],
                "latent_image": ["4", 0],
            },
        },
        "6": {
            "class_type": "VAEDecode",
            "inputs": {"samples": ["5", 0], "vae": ["9", 0]},
        },
        "7": {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "identity-edit", "images": ["6", 0]},
        },
    }


def _queue_prompt(workflow):
    payload = {"prompt": workflow, "client_id": str(uuid.uuid4())}
    response = requests.post(f"{COMFY_URL}/prompt", json=payload, timeout=30)
    if not response.ok:
        raise RuntimeError(f"ComfyUI prompt error {response.status_code}: {response.text[:1000]}")
    return response.json()["prompt_id"]


def _wait_for_history(prompt_id, timeout=900):
    deadline = time.time() + timeout
    while time.time() < deadline:
        response = requests.get(f"{COMFY_URL}/history/{prompt_id}", timeout=30)
        response.raise_for_status()
        history = response.json()
        if prompt_id in history:
            item = history[prompt_id]
            status = item.get("status", {})
            if status.get("completed"):
                return item
            messages = status.get("messages") or []
            for message in messages:
                if isinstance(message, list) and message and message[0] == "execution_error":
                    raise RuntimeError(json.dumps(message[1], indent=2)[:2000])
        time.sleep(1)
    raise RuntimeError("Timed out waiting for ComfyUI generation.")


def _load_output_image(history_item):
    outputs = history_item.get("outputs", {})
    for output in outputs.values():
        for image in output.get("images", []):
            params = {
                "filename": image["filename"],
                "subfolder": image.get("subfolder", ""),
                "type": image.get("type", "output"),
            }
            response = requests.get(f"{COMFY_URL}/view", params=params, timeout=120)
            response.raise_for_status()
            temp_path = Path("/tmp") / f"{uuid.uuid4().hex}.png"
            temp_path.write_bytes(response.content)
            return Image.open(temp_path).convert("RGB")
    raise RuntimeError("ComfyUI completed without returning an image.")


def _duration(prompt, negative_prompt, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler):
    megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024))
    return int(900 + int(steps) * 8 * megapixels)


def _edit_duration(
    input_image,
    prompt,
    negative_prompt,
    width,
    height,
    steps,
    cfg,
    denoise,
    seed,
    randomize_seed,
    sampler,
    scheduler,
):
    megapixels = max(1.0, (int(width) * int(height)) / (1024 * 1024))
    return int(900 + int(steps) * 9 * megapixels)


def _identity_edit_duration(
    input_image,
    prompt,
    ref_boost,
    grounding_px,
    max_megapixels,
    steps,
    cfg,
    seed,
    randomize_seed,
    sampler,
    scheduler,
):
    return int(900 + int(steps) * 10 * max(1.0, float(max_megapixels)))


@spaces.GPU(duration=_duration)
def generate(
    prompt,
    negative_prompt="",
    width=1024,
    height=1024,
    steps=10,
    cfg=1.0,
    seed=0,
    randomize_seed=True,
    sampler="er_sde",
    scheduler="simple",
):
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    try:
        _start_comfyui()
        workflow = _build_workflow(prompt, negative_prompt, width, height, steps, cfg, seed, sampler, scheduler)
        prompt_id = _queue_prompt(workflow)
        history_item = _wait_for_history(prompt_id)
        return _load_output_image(history_item), seed
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=_edit_duration)
def edit_image(
    input_image,
    prompt,
    negative_prompt="",
    width=1024,
    height=1024,
    steps=12,
    cfg=1.2,
    denoise=0.35,
    seed=0,
    randomize_seed=True,
    sampler="er_sde",
    scheduler="simple",
):
    if input_image is None:
        raise gr.Error("Upload an image to edit.")
    if not prompt or not prompt.strip():
        raise gr.Error("Enter an edit prompt.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    try:
        _start_comfyui()
        resized = _resize_for_edit(input_image, width, height)
        input_filename = _upload_image_to_comfy(resized)
        workflow = _build_edit_workflow(
            input_filename,
            prompt,
            negative_prompt,
            steps,
            cfg,
            seed,
            sampler,
            scheduler,
            denoise,
        )
        prompt_id = _queue_prompt(workflow)
        history_item = _wait_for_history(prompt_id)
        return _load_output_image(history_item), seed
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=_identity_edit_duration)
def identity_edit(
    input_image,
    prompt,
    ref_boost=4.0,
    grounding_px=768,
    max_megapixels=1.0,
    steps=10,
    cfg=1.0,
    seed=0,
    randomize_seed=True,
    sampler="euler",
    scheduler="simple",
):
    if input_image is None:
        raise gr.Error("Upload an image to edit.")
    if not prompt or not prompt.strip():
        raise gr.Error("Enter an edit instruction.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    try:
        _start_comfyui()
        source = input_image.convert("RGB")
        width, height = _target_size_from_source(source, max_megapixels)
        input_filename = _upload_image_to_comfy(source)
        workflow = _build_identity_edit_workflow(
            input_filename,
            prompt,
            width,
            height,
            steps,
            cfg,
            seed,
            ref_boost,
            grounding_px,
            sampler,
            scheduler,
        )
        prompt_id = _queue_prompt(workflow)
        history_item = _wait_for_history(prompt_id)
        return _load_output_image(history_item), seed
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


CSS = """
.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
#result-image { min-height: 520px; }
"""

with gr.Blocks(title="Redcraft Krea2", css=CSS) as demo:
    gr.Markdown("# Redcraft Krea2")
    gr.Markdown("ComfyUI-native Redcraft Krea2 generation and image editing.")

    with gr.Tabs():
        with gr.Tab("Generate"):
            with gr.Row():
                with gr.Column(scale=5):
                    prompt = gr.Textbox(label="Prompt", lines=5, placeholder="Describe the image to generate.")
                    negative_prompt = gr.Textbox(label="Negative prompt", lines=2, value="")
                    with gr.Row():
                        width = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                        height = gr.Slider(512, 1536, value=1024, step=64, label="Height")
                    with gr.Row():
                        steps = gr.Slider(1, 30, value=10, step=1, label="Steps")
                        cfg = gr.Slider(0.0, 8.0, value=1.0, step=0.1, label="CFG")
                    with gr.Row():
                        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
                    with gr.Accordion("Sampler", open=False):
                        sampler = gr.Dropdown(
                            ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"],
                            value="er_sde",
                            label="Sampler",
                        )
                        scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler")
                    run = gr.Button("Generate", variant="primary")
                with gr.Column(scale=6):
                    output = gr.Image(label="Result", format="png", elem_id="result-image")

            inputs = [prompt, negative_prompt, width, height, steps, cfg, seed, randomize_seed, sampler, scheduler]
            run.click(generate, inputs, [output, seed])
            prompt.submit(generate, inputs, [output, seed])

        with gr.Tab("Edit Image"):
            with gr.Row():
                with gr.Column(scale=5):
                    edit_input = gr.Image(type="pil", label="Input image")
                    edit_prompt = gr.Textbox(label="Edit prompt", lines=5, placeholder="Describe the edit while preserving identity.")
                    edit_negative_prompt = gr.Textbox(label="Negative prompt", lines=2, value="")
                    with gr.Row():
                        edit_width = gr.Slider(512, 1536, value=1024, step=64, label="Width")
                        edit_height = gr.Slider(512, 1536, value=1024, step=64, label="Height")
                    with gr.Row():
                        edit_steps = gr.Slider(1, 30, value=12, step=1, label="Steps")
                        edit_cfg = gr.Slider(0.0, 8.0, value=1.2, step=0.1, label="CFG")
                    edit_denoise = gr.Slider(
                        0.05,
                        0.8,
                        value=0.35,
                        step=0.05,
                        label="Edit strength",
                        info="Lower values preserve identity and composition more strongly.",
                    )
                    with gr.Row():
                        edit_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        edit_randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
                    with gr.Accordion("Sampler", open=False):
                        edit_sampler = gr.Dropdown(
                            ["er_sde", "euler", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"],
                            value="er_sde",
                            label="Sampler",
                        )
                        edit_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler")
                    edit_run = gr.Button("Edit Image", variant="primary")
                with gr.Column(scale=6):
                    edit_output = gr.Image(label="Edited image", format="png", elem_id="result-image")

            edit_inputs = [
                edit_input,
                edit_prompt,
                edit_negative_prompt,
                edit_width,
                edit_height,
                edit_steps,
                edit_cfg,
                edit_denoise,
                edit_seed,
                edit_randomize_seed,
                edit_sampler,
                edit_scheduler,
            ]
            edit_run.click(edit_image, edit_inputs, [edit_output, edit_seed])
            edit_prompt.submit(edit_image, edit_inputs, [edit_output, edit_seed])

        with gr.Tab("Identity Edit"):
            gr.Markdown(
                "Instruction-based, identity-preserving editing using the community LoRA "
                "[`conradlocke/krea2-identity-edit`](https://huggingface.co/conradlocke/krea2-identity-edit) "
                "on the Redcraft Krea2 checkpoint, via the "
                "[ComfyUI-Krea2Edit](https://github.com/lbouaraba/comfyui-krea2edit) node pack. "
                "Give it an image and a plain-language instruction; it edits while preserving what you "
                "didn't ask to change, including the person."
            )
            with gr.Row():
                with gr.Column(scale=5):
                    id_input = gr.Image(type="pil", label="Source image")
                    id_prompt = gr.Textbox(
                        label="Edit instruction",
                        lines=3,
                        placeholder="e.g. create a photo of this person at a night market",
                    )
                    id_ref_boost = gr.Slider(
                        0.0,
                        10.0,
                        value=4.0,
                        step=0.5,
                        label="Likeness (ref_boost)",
                        info="How hard the edit holds the reference. 1 = off, 4 = strong likeness (recommended), 8+ over-copies.",
                    )
                    with gr.Accordion("Advanced settings", open=False):
                        id_grounding_px = gr.Slider(
                            384,
                            1536,
                            value=768,
                            step=64,
                            label="Grounding resolution (px)",
                            info="Lower = stronger edit adherence. Higher = stronger identity likeness. Trained range 384-768; 1024+ often still works.",
                        )
                        id_max_mp = gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Output size (megapixels)")
                        with gr.Row():
                            id_steps = gr.Slider(4, 28, value=10, step=1, label="Steps")
                            id_cfg = gr.Slider(
                                0.0,
                                5.0,
                                value=1.0,
                                step=0.5,
                                label="CFG",
                                info="Turbo convention: 1.0 = guidance off. Raise (e.g. 3) with more steps for removals/large edits.",
                            )
                        with gr.Row():
                            id_seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                            id_randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
                        with gr.Row():
                            id_sampler = gr.Dropdown(
                                ["euler", "er_sde", "euler_ancestral", "dpmpp_2m", "dpmpp_sde"],
                                value="euler",
                                label="Sampler",
                            )
                            id_scheduler = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="simple", label="Scheduler")
                    id_run = gr.Button("Edit Identity", variant="primary")
                with gr.Column(scale=6):
                    id_output = gr.Image(label="Edited image", format="png", elem_id="result-image")

            id_inputs = [
                id_input,
                id_prompt,
                id_ref_boost,
                id_grounding_px,
                id_max_mp,
                id_steps,
                id_cfg,
                id_seed,
                id_randomize_seed,
                id_sampler,
                id_scheduler,
            ]
            id_run.click(identity_edit, id_inputs, [id_output, id_seed])
            id_prompt.submit(identity_edit, id_inputs, [id_output, id_seed])

            gr.Examples(
                examples=[
                    ["examples/woman.jpg", "create a photo of this person at a busy night market at night"],
                    ["examples/businessman_suit.jpg", "change the suit jacket to a red leather jacket"],
                    ["examples/man_beach.jpg", "make it a vintage film photo with warm golden-hour light"],
                ],
                inputs=[id_input, id_prompt],
                label="Examples",
            )


_ensure_comfyui()

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