app.py

"""Krea 2 Identity Edit — instruction-based, identity-preserving image editing.

This Space runs the community LoRA `conradlocke/krea2-identity-edit` on top of
`krea/Krea-2-Turbo` (the distilled 8-step checkpoint). The LoRA is trained with a
*dual conditioning* recipe that stock text-to-image `Krea2Pipeline` does not
provide, so this app reproduces the two custom pieces from the reference
ComfyUI-Krea2Edit node pack (https://github.com/lbouaraba/comfyui-krea2edit) in
plain diffusers / PyTorch:

  1. Krea2EditGroundedEncode — the instruction is encoded *together with the
     source image* through the Qwen3-VL text encoder (vision tokens inserted via
     the image-grounded chat template), and the 12 selected decoder layers are
     tapped, exactly like the text-only Krea 2 path but with the image grounding
     the semantics ("the man on the left", "the sign in the back").

  2. Krea2EditModelPatch — the VAE-encoded SOURCE latent is prepended to the
     transformer sequence as a block of *clean* tokens, distinguished from the
     noisy target purely by the 3-axis RoPE frame index (source frame = 1,
     target frame = 0, h/w aligned). The sequence becomes
     [text | source(frame=1) | target(frame=0)] and only the target tokens are
     kept as the velocity prediction — mirroring ai-toolkit's
     `predict_velocity_edit`.

Everything runs on ZeroGPU: modules go on CUDA at module scope, inference is
wrapped in @spaces.GPU, no torch.compile, no CPU offload.
"""

import os

# Disable Gradio server-side rendering: SSR + the gradio_client / ZeroGPU
# interaction has been observed to surface as opaque CUDA allocator asserts.
os.environ.setdefault("GRADIO_SSR_MODE", "False")

import random

import numpy as np
import spaces
import torch
from torch.nn.attention import sdpa_kernel, SDPBackend
import gradio as gr
from PIL import Image

from diffusers import Krea2Pipeline
from diffusers.pipelines.krea2.pipeline_krea2 import retrieve_timesteps
from transformers import AutoProcessor

# --------------------------------------------------------------------------------------
# Constants
# --------------------------------------------------------------------------------------
BASE_MODEL = "krea/Krea-2-Turbo"
LORA_REPO = "conradlocke/krea2-identity-edit"
LORA_WEIGHT = "krea2_identity_edit_v1_2.safetensors"
DTYPE = torch.bfloat16
MAX_SEED = np.iinfo(np.int32).max

# Turbo recipe for "most edits" (add / recolor / restyle / re-stage), per the LoRA card:
#   Turbo, ~10 steps, CFG 1.0 (== guidance disabled in the Krea convention).
#   v1.2: 8 = stronger composition/instruction adherence, 12 = more face detail.
DEFAULT_STEPS = 10
DEFAULT_GUIDANCE = 0.0          # Krea convention: 0.0 disables guidance
DEFAULT_GROUNDING_PX = 768      # trained dial 384-768; higher often still works
DEFAULT_LORA_SCALE = 1.0
DEFAULT_REF_BOOST = 4.0         # v1.2 likeness dial: 1.0 = off, ~4 = strong likeness
                                # (the release workflow's preset); >10 over-copies.
MAX_MEGAPIXELS = 1.0           # card allows <=2MP; we cap at 1MP because the edit
                               # path prepends the source latent (~2x image tokens).
                               # Inference runs under torch.inference_mode(), so no
                               # autograd buffers are held.

# The image-grounded instruction template from ComfyUI-Krea2Edit. The system
# prefix is byte-identical to the diffusers Krea 2 text template; the difference
# is the <|vision_start|><|image_pad|><|vision_end|> block inserted before the
# instruction so the VLM grounds the edit on the source image.
GROUNDED_TEMPLATE = (
    "<|im_start|>system\nDescribe the image by detailing the color, shape, size, "
    "texture, quantity, text, spatial relationships of the objects and background:"
    "<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
    "{}<|im_end|>\n<|im_start|>assistant\n"
)

# --------------------------------------------------------------------------------------
# Load the pipeline (module scope, eager CUDA placement — required for ZeroGPU)
# --------------------------------------------------------------------------------------
pipe = Krea2Pipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE)


def _convert_lora_keys(state_dict):
    """Convert the ai-toolkit / ComfyUI Krea2 LoRA naming to the diffusers
    `Krea2Transformer2DModel` module names so `load_lora_adapter` matches.

    ai-toolkit keys look like `diffusion_model.blocks.0.attn.wq.lora_A.weight`;
    diffusers expects `transformer_blocks.0.attn.to_q.lora_A.weight` (and
    `txtfusion` -> `text_fusion`, `mlp` -> `ff`, `wo` -> `to_out.0`, etc.).
    """
    attn_map = {
        ".attn.wq.": ".attn.to_q.",
        ".attn.wk.": ".attn.to_k.",
        ".attn.wv.": ".attn.to_v.",
        ".attn.wo.": ".attn.to_out.0.",
        ".attn.gate.": ".attn.to_gate.",
        ".mlp.gate.": ".ff.gate.",
        ".mlp.up.": ".ff.up.",
        ".mlp.down.": ".ff.down.",
    }
    out = {}
    for k, v in state_dict.items():
        nk = k
        if nk.startswith("diffusion_model."):
            nk = nk[len("diffusion_model."):]
        nk = nk.replace("blocks.", "transformer_blocks.", 1) if nk.startswith("blocks.") else nk
        nk = nk.replace("txtfusion.", "text_fusion.", 1) if nk.startswith("txtfusion.") else nk
        for a, b in attn_map.items():
            if a in nk:
                nk = nk.replace(a, b)
        out[nk] = v
    return out


# Load the identity-edit LoRA onto the transformer (Krea 2 LoRAs load through the
# transformer's adapter API). The weights ship in ai-toolkit naming, so convert
# the keys to the diffusers convention first.
from safetensors.torch import load_file as _load_safetensors
from huggingface_hub import hf_hub_download as _hf_hub_download

_lora_path = _hf_hub_download(LORA_REPO, LORA_WEIGHT)
_lora_sd = _convert_lora_keys(_load_safetensors(_lora_path))
# Keys are already in diffusers module-path form (transformer_blocks.N.attn.to_q...),
# so bypass the default prefix="transformer" filter (which would drop every key
# because "transformer_blocks." does not start with "transformer.").
pipe.transformer.load_lora_adapter(_lora_sd, adapter_name="default", prefix=None)
pipe.transformer.set_adapters("default", weights=DEFAULT_LORA_SCALE)
# Fuse the LoRA into the base weights at strength 1.0 (the card's recommended
# setting) and unload the adapter. This bakes the delta into the linear weights so
# inference runs through plain nn.Linear rather than peft's LoRA layer wrapper —
# faster, and it sidesteps a ZeroGPU caching-allocator assert triggered inside the
# peft LoRA forward. The adjustable lora_scale dial is therefore fixed at 1.0.
pipe.transformer.fuse_lora(lora_scale=DEFAULT_LORA_SCALE)
pipe.transformer.unload_lora()

# --- ref_boost kernel compatibility -------------------------------------------------
# Krea2 attention is GQA (48 query heads / 12 kv heads). With a dense float
# attn_mask (the ref_boost bias) every fused SDPA kernel refuses GQA inputs
# ("dense input requires q/k/v same num_heads") -> "No available kernel".
# The reference ComfyUI node sidesteps this by always repeat_interleaving K/V to
# the full head count before attention. Do the same here, but only when a mask
# is present: mask=None keeps the stock GQA fast path (identical to upstream).
from diffusers.models.transformers.transformer_krea2 import (  # noqa: E402
    Krea2AttnProcessor as _K2Proc,
)
from diffusers.models.attention_dispatch import dispatch_attention_fn as _dispatch  # noqa: E402
from diffusers.models.embeddings import apply_rotary_emb as _apply_rope  # noqa: E402


class _MaskCompatProcessor(_K2Proc):
    """Krea2AttnProcessor + dense-mask support under GQA (kv repeated to full heads)."""

    def __call__(self, attn, hidden_states, attention_mask=None, image_rotary_emb=None):
        if attention_mask is None:  # stock fast path, byte-identical behavior
            return super().__call__(attn, hidden_states, attention_mask, image_rotary_emb)
        query = attn.to_q(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
        key = attn.to_k(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
        value = attn.to_v(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
        gate = attn.to_gate(hidden_states)
        query = attn.norm_q(query)
        key = attn.norm_k(key)
        if image_rotary_emb is not None:
            query = _apply_rope(query, image_rotary_emb, sequence_dim=1)
            key = _apply_rope(key, image_rotary_emb, sequence_dim=1)
        rep = attn.num_heads // attn.num_kv_heads
        if rep > 1:  # materialize GQA so the mem-efficient kernel accepts the mask
            key = key.repeat_interleave(rep, dim=2)      # (B, L, H, D) layout
            value = value.repeat_interleave(rep, dim=2)
        hidden_states = _dispatch(
            query, key, value,
            attn_mask=attention_mask,
            enable_gqa=False,
            backend=self._attention_backend,
            parallel_config=self._parallel_config,
        )
        hidden_states = hidden_states.flatten(2, 3)
        hidden_states = hidden_states * torch.sigmoid(gate)
        return attn.to_out[0](hidden_states)


for _blk in pipe.transformer.transformer_blocks:
    _blk.attn.set_processor(_MaskCompatProcessor())

pipe.to("cuda")

# A Qwen3-VL processor for the grounded (image + text) encode. The Krea 2 repo
# ships only a text tokenizer; the vision-side preprocessing (image mean/std,
# 16px patch, spatial-merge=2) comes from the Qwen3-VL processor. We reuse the
# Krea 2 tokenizer's special tokens by aligning the template above with the
# processor's <|image_pad|> expansion.
try:
    processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-4B-Instruct")
except Exception as e:  # pragma: no cover - surfaced in logs if it happens
    print(f"[warn] could not load Qwen3-VL processor, falling back: {e}")
    processor = None

# VAE latent normalization stats (the transformer works in normalized latent
# space; randn target latents already live there, so the source must be
# normalized the same way as the pipeline: (z - mean) / std).
_LATENTS_MEAN = (
    torch.tensor(pipe.vae.config.latents_mean).view(1, pipe.vae.config.z_dim, 1, 1, 1)
)
_LATENTS_STD = (
    torch.tensor(pipe.vae.config.latents_std).view(1, pipe.vae.config.z_dim, 1, 1, 1)
)


# --------------------------------------------------------------------------------------
# Grounded instruction encoding (semantic path)
# --------------------------------------------------------------------------------------
def _grounded_encode(instruction: str, source: Image.Image, grounding_px: int):
    """Encode the instruction grounded on the source image through Qwen3-VL.

    Returns (prompt_embeds, prompt_embeds_mask) shaped like the diffusers
    Krea 2 text conditioning: (1, seq, num_text_layers, text_hidden_dim) and
    (1, seq).
    """
    device = pipe._execution_device
    select_layers = pipe.text_encoder_select_layers
    prefix_idx = pipe.prompt_template_encode_start_idx  # 34: drop the system prefix

    # Cap the longest side fed to the VLM (the LoRA trained with 384-768px jitter).
    img = source.convert("RGB")
    if grounding_px and max(img.size) > grounding_px:
        s = grounding_px / max(img.size)
        img = img.resize(
            (max(16, round(img.size[0] * s)), max(16, round(img.size[1] * s))),
            Image.LANCZOS,
        )

    text = GROUNDED_TEMPLATE.format(instruction or "")

    inputs = processor(
        text=[text],
        images=[img],
        padding=True,
        return_tensors="pt",
    ).to(device)

    te_kwargs = dict(
        input_ids=inputs["input_ids"],
        attention_mask=inputs.get("attention_mask"),
        pixel_values=inputs.get("pixel_values"),
        image_grid_thw=inputs.get("image_grid_thw"),
        output_hidden_states=True,
    )
    # Qwen3-VL needs mm_token_type_ids (returned by the processor) to compute
    # multimodal M-RoPE positions when image tokens are present.
    if inputs.get("mm_token_type_ids") is not None:
        te_kwargs["mm_token_type_ids"] = inputs["mm_token_type_ids"]
    outputs = pipe.text_encoder(**te_kwargs)
    hidden_states = torch.stack(
        [outputs.hidden_states[i] for i in select_layers], dim=2
    )  # (1, seq, num_layers, dim)

    attention_mask = inputs.get("attention_mask")
    if attention_mask is None:
        attention_mask = torch.ones(
            hidden_states.shape[:2], device=device, dtype=torch.bool
        )
    else:
        attention_mask = attention_mask.bool()

    # Drop the system-prefix tokens (identical prefix to the text-only path).
    hidden_states = hidden_states[:, prefix_idx:]
    attention_mask = attention_mask[:, prefix_idx:]
    return hidden_states.to(DTYPE), attention_mask


def _text_only_encode(instruction: str):
    """Fallback text-only encode if the processor is unavailable."""
    return pipe.encode_prompt(prompt=instruction, device=pipe._execution_device)


# --------------------------------------------------------------------------------------
# Source-preservation forward (appearance path)
# --------------------------------------------------------------------------------------
def _encode_source_latent(source: Image.Image, height: int, width: int):
    """VAE-encode the source image to a packed, normalized latent block matching
    the target grid, ready to prepend to the transformer sequence."""
    device = pipe._execution_device
    # Preprocess to the target resolution (training pairs are same-size; the card
    # says match output AR to the source, which we enforce upstream).
    px = pipe.image_processor.preprocess(source.convert("RGB"), height=height, width=width)
    px = px.unsqueeze(2).to(device=device, dtype=pipe.vae.dtype)  # (B,C,1,H,W)

    latent = pipe.vae.encode(px).latent_dist.mode()  # (B, z, 1, lh, lw), unnormalized
    mean = _LATENTS_MEAN.to(latent.device, latent.dtype)
    std = _LATENTS_STD.to(latent.device, latent.dtype)
    latent = (latent - mean) / std               # normalized latent space (matches pipeline)
    latent = latent[:, :, 0]                      # (B, z, lh, lw)

    b, c, lh, lw = latent.shape
    packed = pipe._pack_latents(latent, b, c, lh, lw)  # (B, lh*lw/p^2, c*p*p)
    return packed.to(DTYPE)


def _edit_position_ids(text_seq_len, grid_h, grid_w, n_src, device):
    """Build (text + n_src*grid + grid, 3) rotary coords:
    text @ (0,0,0); each source block @ frame=(i+1) with (h,w); target @ frame=0.
    """
    text_ids = torch.zeros(text_seq_len, 3, device=device)

    def _img_ids(frame):
        ids = torch.zeros(grid_h, grid_w, 3, device=device)
        ids[..., 0] = frame
        ids[..., 1] = torch.arange(grid_h, device=device)[:, None]
        ids[..., 2] = torch.arange(grid_w, device=device)[None, :]
        return ids.reshape(grid_h * grid_w, 3)

    blocks = [text_ids]
    blocks += [_img_ids(i + 1) for i in range(n_src)]  # sources frame=1..N
    blocks += [_img_ids(0)]                              # target frame=0
    return torch.cat(blocks, dim=0)


def _ref_boost_bias(text_len, src_len, tgt_len, ref_boost, device, dtype):
    """Additive attention-logit bias implementing the v1.2 `ref_boost` fidelity
    dial (ComfyUI-Krea2Edit's _ref_attn_bias, single-ref case): target queries
    get log(ref_boost) added on the source-block keys — equivalent to
    multiplying the refs' post-softmax attention weight before renormalization.
    1.0 = off (return None and keep the maskless fast path). The float bias is
    handled chunk-wise by the mem-efficient SDPA backend, so the full score
    matrix is never materialized (the OOM note above concerns *boolean* masks,
    which force the math path)."""
    import math as _math
    if ref_boost == 1.0:
        return None
    L = text_len + src_len + tgt_len
    bias = torch.zeros(1, 1, L, L, device=device, dtype=dtype)
    rows0 = text_len + src_len                       # target rows
    bias[:, :, rows0:, text_len:rows0] = _math.log(max(float(ref_boost), 1e-4))
    return bias


def _edit_transformer_forward(latents, src_packed, prompt_embeds, prompt_mask,
                              timestep, position_ids, ref_boost=1.0):
    """Run the Krea 2 transformer with the source latent block prepended, keeping
    only the target tokens out. Reproduces ComfyUI-Krea2Edit's krea2_edit_forward
    against the diffusers Krea2Transformer2DModel (img_in == m.first,
    transformer_blocks == m.blocks, text_fusion/txt_in == m.txtfusion/m.txtmlp,
    rotary_emb == m.pe_embedder, final_layer == m.last)."""
    m = pipe.transformer
    combined_img = torch.cat([src_packed, latents], dim=1)  # [source | target] packed

    temb = m.time_embed(timestep, dtype=latents.dtype)
    temb_mod = m.time_mod_proj(torch.nn.functional.gelu(temb, approximate="tanh"))

    # Text fusion + projection (attention mask over text only; all image tokens valid).
    text_attn_mask = prompt_mask[:, None, None, :] if prompt_mask is not None else None
    enc = m.text_fusion(prompt_embeds, attention_mask=text_attn_mask)
    enc = m.txt_in(enc)

    img = m.img_in(combined_img)
    hidden = torch.cat([enc, img], dim=1)  # [text | source | target]

    image_rotary_emb = m.rotary_emb(position_ids)

    # Single prompt, batch=1: the grounded encode uses padding=True on one
    # sequence, so there is no padding and every token is valid. With
    # ref_boost=1.0 the mask stays None and SDPA picks the flash / mem-efficient
    # kernels (a dense *boolean* mask would force the score-materializing math
    # path and OOM given the doubled token count). If padding is ever introduced
    # (batched or CFG with unequal lengths), trim to the shorter length instead
    # of masking.
    attention_mask = _ref_boost_bias(
        enc.shape[1], src_packed.shape[1], latents.shape[1],
        ref_boost, hidden.device, hidden.dtype,
    )

    # Prefer the flash / mem-efficient SDPA kernels (they never materialize the
    # full seq x seq score matrix, which the MATH kernel would OOM on given the
    # doubled token count from prepending the source latent). With a float
    # ref_boost bias, flash declines the mask and SDPA falls through to the
    # mem-efficient kernel, which applies the bias chunk-wise.
    with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
        for block in m.transformer_blocks:
            hidden = block(hidden, temb_mod, image_rotary_emb, attention_mask)

    text_seq_len = enc.shape[1]
    tgt_len = latents.shape[1]
    hidden = hidden[:, text_seq_len:]              # drop text -> [source | target]
    hidden = hidden[:, -tgt_len:]                  # keep target tokens only
    return m.final_layer(hidden, temb)


# --------------------------------------------------------------------------------------
# Sizing helpers
# --------------------------------------------------------------------------------------
def _target_size(source: Image.Image):
    """Match the output AR to the source (card requirement) and cap at ~2MP,
    snapping each side to a multiple of vae_scale_factor * patch_size."""
    multiple = pipe.vae_scale_factor * pipe.patch_size  # 8 * 2 = 16
    w, h = source.size
    mp = (w * h) / 1e6
    if mp > MAX_MEGAPIXELS:
        s = (MAX_MEGAPIXELS / mp) ** 0.5
        w, h = round(w * s), round(h * s)
    w = max(multiple, (w // multiple) * multiple)
    h = max(multiple, (h // multiple) * multiple)
    return h, w


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
@spaces.GPU
@torch.inference_mode()
def edit(
    source_image,
    instruction,
    ref_boost=DEFAULT_REF_BOOST,
    grounding_px=DEFAULT_GROUNDING_PX,
    lora_scale=DEFAULT_LORA_SCALE,
    steps=DEFAULT_STEPS,
    guidance_scale=DEFAULT_GUIDANCE,
    seed=0,
    randomize_seed=True,
    progress=gr.Progress(track_tqdm=True),
):
    if source_image is None:
        raise gr.Error("Please upload a source image to edit.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please describe the edit you want (e.g. 'put this person at a night market').")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    device = pipe._execution_device
    source = source_image if isinstance(source_image, Image.Image) else Image.fromarray(source_image)
    source = source.convert("RGB")

    # The identity-edit LoRA is fused into the base weights at load time (strength
    # 1.0, the card's recommended setting), so lora_scale is fixed and not re-applied
    # per call.

    height, width = _target_size(source)

    # --- Semantic path: grounded instruction encode -----------------------------------
    if processor is not None:
        prompt_embeds, prompt_mask = _grounded_encode(instruction.strip(), source, int(grounding_px))
    else:
        prompt_embeds, prompt_mask = _text_only_encode(instruction.strip())

    do_cfg = float(guidance_scale) > 0
    if do_cfg:
        # At CFG > 1 the card says ground the negative too: empty prompt, same image.
        if processor is not None:
            neg_embeds, neg_mask = _grounded_encode("", source, int(grounding_px))
        else:
            neg_embeds, neg_mask = _text_only_encode("")

    # --- Appearance path: encode + pack the source latent -----------------------------
    src_packed = _encode_source_latent(source, height, width)

    # --- Prepare noisy target latents -------------------------------------------------
    num_channels_latents = pipe.transformer.config.in_channels // (pipe.patch_size ** 2)
    generator = torch.Generator(device=device).manual_seed(seed)
    latents = pipe.prepare_latents(
        1, num_channels_latents, height, width, DTYPE, device, generator, None
    )

    grid_h = height // (pipe.vae_scale_factor * pipe.patch_size)
    grid_w = width // (pipe.vae_scale_factor * pipe.patch_size)
    position_ids = _edit_position_ids(prompt_embeds.shape[1], grid_h, grid_w, 1, device)

    # --- Timesteps (distilled schedule: fixed mu = 1.15) ------------------------------
    sigmas = np.linspace(1.0, 1 / int(steps), int(steps))
    timesteps, num_steps = retrieve_timesteps(
        pipe.scheduler, int(steps), device, sigmas=sigmas, mu=1.15
    )

    # --- Denoising loop ---------------------------------------------------------------
    pipe.scheduler.set_begin_index(0)
    for t in progress.tqdm(timesteps, desc="Editing"):
        timestep = (t / pipe.scheduler.config.num_train_timesteps).expand(latents.shape[0]).to(latents.dtype)

        noise_pred = _edit_transformer_forward(
            latents, src_packed, prompt_embeds, prompt_mask, timestep, position_ids,
            ref_boost=float(ref_boost),
        )
        if do_cfg:
            neg_position_ids = _edit_position_ids(neg_embeds.shape[1], grid_h, grid_w, 1, device)
            neg_pred = _edit_transformer_forward(
                latents, src_packed, neg_embeds, neg_mask, timestep, neg_position_ids,
                ref_boost=float(ref_boost),
            )
            noise_pred = noise_pred + float(guidance_scale) * (noise_pred - neg_pred)

        latents = pipe.scheduler.step(noise_pred, t, latents, return_dict=False)[0]

    # --- Decode -----------------------------------------------------------------------
    latents = pipe._unpack_latents(latents, height, width).to(pipe.vae.dtype)
    mean = _LATENTS_MEAN.to(latents.device, latents.dtype)
    std = _LATENTS_STD.to(latents.device, latents.dtype)
    latents = latents * std + mean               # denorm matches pipeline: normalized * std + mean
    image = pipe.vae.decode(latents, return_dict=False)[0][:, :, 0]
    image = pipe.image_processor.postprocess(image, output_type="pil")[0]

    return image, seed


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
"""

DESCRIPTION = """
# 🧬 Krea 2 Identity Edit

Instruction-based, **identity-preserving** image editing on
[Krea 2 Turbo](https://huggingface.co/krea/Krea-2-Turbo) with the community LoRA
[`conradlocke/krea2-identity-edit`](https://huggingface.co/conradlocke/krea2-identity-edit).

**Now running v1.2** — better face likeness, head/face swap, outpainting, inpainting,
try-on, and the new **Likeness dial** below (`ref_boost`, ships at 4).
[ComfyUI nodes](https://github.com/lbouaraba/comfyui-krea2edit) for local use.

A solo hobby project, free and open. Optional tip jar for GPU compute: [ko-fi.com/conradlocke](https://ko-fi.com/conradlocke) ☕
"""

with gr.Blocks(css=CSS, title="Krea 2 Identity Edit") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)

        with gr.Row(equal_height=True):
            with gr.Column():
                source_image = gr.Image(label="Source image", type="pil", height=420)
                instruction = gr.Textbox(
                    label="Edit instruction",
                    placeholder="e.g. create a photo of this person at a night market",
                    lines=2,
                )
                run_button = gr.Button("Edit", variant="primary", size="lg")

                ref_boost = gr.Slider(
                    label="Likeness (ref_boost)",
                    minimum=0.0, maximum=10.0, step=0.5, value=DEFAULT_REF_BOOST,
                    info="v1.2 fidelity dial: how hard the edit holds the reference. "
                         "1 = off (v1.1 behavior), 4 = strong likeness (recommended). "
                         "8+ over-copies: stylized/comic edits and removals start failing.",
                )

                with gr.Accordion("Advanced settings", open=False):
                    grounding_px = gr.Slider(
                        label="Grounding resolution (px)",
                        minimum=512, maximum=1536, step=64, value=DEFAULT_GROUNDING_PX,
                        info="Lower = stronger edit adherence; higher = stronger identity/likeness. Try 1024+ for people.",
                    )
                    lora_scale = gr.Slider(
                        label="LoRA strength (fused at 1.0)",
                        minimum=0.0, maximum=1.5, step=0.05, value=DEFAULT_LORA_SCALE,
                        interactive=False,
                        info="The identity-edit LoRA is fused into the base weights at the "
                             "card-recommended strength 1.0.",
                    )
                    steps = gr.Slider(
                        label="Steps",
                        minimum=4, maximum=28, step=1, value=DEFAULT_STEPS,
                        info="Turbo default is 8. For removals/large deletions try more steps + guidance > 0.",
                    )
                    guidance_scale = gr.Slider(
                        label="Guidance (0 = off, Turbo default)",
                        minimum=0.0, maximum=5.0, step=0.5, value=DEFAULT_GUIDANCE,
                        info="Krea convention: 0 disables guidance. Raise (e.g. 3) for removals/large edits.",
                    )
                    with gr.Row():
                        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Column():
                result = gr.Image(label="Edited image", height=420)
                used_seed = gr.Number(label="Seed used", interactive=False)

        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=[source_image, instruction],
            outputs=[result, used_seed],
            fn=edit,
            cache_examples=True,
            cache_mode="lazy",
        )

    gr.on(
        triggers=[run_button.click, instruction.submit],
        fn=edit,
        inputs=[source_image, instruction, ref_boost, grounding_px, lora_scale, steps, guidance_scale, seed, randomize_seed],
        outputs=[result, used_seed],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, show_error=True)