Revert to original (undo accidental overwrite)

app.py

@@ -1,195 +1,510 @@
+"""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 diffusers import DiffusionPipeline
+from PIL import Image
 
-BASE_MODEL = "krea/Krea-2-Turbo"
-CUSTOM_PIPELINE = "ostris/Krea2OstrisEdit"
+from diffusers import Krea2Pipeline
+from diffusers.pipelines.krea2.pipeline_krea2 import retrieve_timesteps
+from transformers import AutoProcessor
 
-# Instruction-based, identity-preserving edit LoRA. This IS the edit brain for
-# the pipeline's dual conditioning (in-context VAE tokens + Qwen3-VL grounding),
-# so it replaces the Ostris style-reference LoRA rather than stacking on it.
+# --------------------------------------------------------------------------------------
+# Constants
+# --------------------------------------------------------------------------------------
+BASE_MODEL = "krea/Krea-2-Turbo"
 LORA_REPO = "conradlocke/krea2-identity-edit"
-LORA_WEIGHT = "krea2_identity_edit_v1_1.safetensors"  # v1.1 recommended (full rank)
-
+LORA_WEIGHT = "krea2_identity_edit_v1_1.safetensors"
 DTYPE = torch.bfloat16
-MAX_SEED = 2**31 - 1
-TARGET_PX = 1024  # ~1MP target; keep well under the LoRA's ~1.5-2MP ceiling
+MAX_SEED = np.iinfo(np.int32).max
 
+# Turbo recipe for "most edits" (add / recolor / restyle / re-stage), per the LoRA card:
+#   Turbo, 8 steps, CFG 1.0 (== guidance disabled in the Krea convention).
+DEFAULT_STEPS = 8
+DEFAULT_GUIDANCE = 0.0          # Krea convention: 0.0 disables guidance
+DEFAULT_GROUNDING_PX = 768      # trained dial 512-1536; 768 balanced, 1024+ for people
+DEFAULT_LORA_SCALE = 1.0
+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.
 
-pipe = DiffusionPipeline.from_pretrained(
-    BASE_MODEL,
-    custom_pipeline=CUSTOM_PIPELINE,
-    torch_dtype=DTYPE,
-    trust_remote_code=True,
+# 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()
+
 pipe.to("cuda")
-pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHT)
-
-
-def _match_source_size(image, base=TARGET_PX, mult=16, lo=512, hi=1536):
-    """Output AR must match the source (training pairs are same-size); target
-    ~base*base area, preserve aspect ratio, snap to multiples of 16."""
-    w, h = image.size
-    ar = w / h
-    th = (base * base / ar) ** 0.5
-    tw = th * ar
-    tw = max(lo, min(hi, round(tw / mult) * mult))
-    th = max(lo, min(hi, round(th / mult) * mult))
-    return int(tw), int(th)
-
-
-@spaces.GPU(duration=90, size="large")
-def generate(
-    prompt,
-    reference_image,
-    use_reference,
-    lora_scale=1.0,
-    steps=8,
-    guidance=1.0,
+
+# 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 _edit_transformer_forward(latents, src_packed, prompt_embeds, prompt_mask,
+                              timestep, position_ids):
+    """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. Passing
+    # attention_mask=None lets SDPA pick 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 = None
+
+    # 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 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,
+    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 not prompt or not prompt.strip():
-        raise gr.Error("Enter an edit instruction, e.g. 'put a red party hat on it'.")
-    if use_reference and reference_image is None:
-        raise gr.Error("Upload a photo to edit, or disable the reference option.")
+    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 or seed is None:
+    if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     seed = int(seed)
 
-    lora_scale = 1.0 if lora_scale is None else float(lora_scale)
-    steps = 8 if steps is None else int(steps)
-    guidance = 1.0 if guidance is None else float(guidance)
+    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)
 
-    # Match output size/AR to the source image; fall back to square otherwise.
-    if use_reference and reference_image is not None:
-        width, height = _match_source_size(reference_image)
+    # --- Semantic path: grounded instruction encode -----------------------------------
+    if processor is not None:
+        prompt_embeds, prompt_mask = _grounded_encode(instruction.strip(), source, int(grounding_px))
     else:
-        width = height = TARGET_PX
-
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    pipe_kwargs = dict(
-        prompt=prompt.strip(),
-        num_inference_steps=steps,
-        guidance_scale=guidance,
-        width=width,
-        height=height,
-        generator=generator,
-        attention_kwargs={"scale": lora_scale},
+        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
     )
-    if use_reference and reference_image is not None:
-        pipe_kwargs["image"] = reference_image
 
-    image = pipe(**pipe_kwargs).images[0]
+    # --- 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
+        )
+        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
+            )
+            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 = """
-#page { max-width: 1100px; margin: 0 auto; padding: 4px 8px 32px; }
-#header { padding: 24px 4px 18px; border-bottom: 1px solid #e5e5e5; margin-bottom: 20px; }
-#header h1 { font-size: 32px; font-weight: 700; margin: 0 0 6px; letter-spacing: -0.02em; }
-#header .subtitle { font-size: 15px; color: #666; margin: 0; max-width: 72ch; line-height: 1.5; }
-#header .links { margin-top: 12px; display: flex; gap: 16px; }
-#header .links a {
-    font-size: 13px; color: #888; text-decoration: none;
-    border: 1px solid #ddd; border-radius: 6px; padding: 4px 10px;
-}
-#header .links a:hover { color: #333; border-color: #aaa; }
-footer { display: none !important; }
+#col-container { max-width: 1100px; margin: 0 auto; }
 """
 
-HEADER = """
-<div id="header">
-  <h1>Krea 2 Identity Edit ✏️</h1>
-  <p class="subtitle">
-    Instruction-based, <b>identity-preserving</b> image editing. Upload a photo and give a
-    plain-language instruction — "put this person at a night market", "add a red party hat",
-    "make it snow" — and the
-    <a href="https://huggingface.co/conradlocke/krea2-identity-edit" target="_blank">Krea2 Identity Edit LoRA</a>
-    edits it while keeping the face, outfit, and untouched parts of the scene fixed. Built on
-    <a href="https://huggingface.co/krea/Krea-2-Turbo" target="_blank">Krea 2 Turbo</a>
-    via the
-    <a href="https://huggingface.co/ostris/Krea2OstrisEdit" target="_blank">Krea2OstrisEdit</a>
-    pipeline.
-  </p>
-  <div class="links">
-    <a href="https://huggingface.co/conradlocke/krea2-identity-edit" target="_blank">Edit LoRA ↗</a>
-    <a href="https://huggingface.co/ostris/Krea2OstrisEdit" target="_blank">Pipeline ↗</a>
-    <a href="https://huggingface.co/krea/Krea-2-Turbo" target="_blank">Base model ↗</a>
-  </div>
-</div>
+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).
 """
 
-with gr.Blocks(title="Krea 2 Identity Edit") as demo:
-    with gr.Column(elem_id="page"):
-        gr.HTML(HEADER)
-        with gr.Row(equal_height=False):
-            with gr.Column(scale=1):
-                prompt = gr.Textbox(
+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,
-                    placeholder="e.g. 'put a red party hat on it' or 'place this person on a beach at sunset'",
-                )
-                use_reference = gr.Checkbox(
-                    value=True,
-                    label="Use uploaded photo",
-                    info="On: edit your photo. Off: plain text-to-image from the instruction.",
-                )
-                reference_image = gr.Image(
-                    label="Photo to edit",
-                    type="pil",
-                    height=300,
                 )
-                generate_btn = gr.Button("Edit", variant="primary", size="lg")
-                with gr.Accordion("Advanced", open=False):
+                run_button = gr.Button("Edit", variant="primary", size="lg")
+
+                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(
-                        0.0, 1.5, value=1.0, step=0.05,
-                        label="Edit LoRA strength",
-                        info="1.0 recommended",
+                        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.",
                     )
-                    steps = gr.Slider(1, 30, value=8, step=1, label="Steps",
-                                      info="8–12 for most edits; ~20 for large removals")
-                    guidance = gr.Slider(
-                        0.0, 6.0, value=1.0, step=0.1,
-                        label="Guidance (CFG)",
-                        info="1.0 for most edits; ~3.0 for deletions",
+                    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(0, MAX_SEED, value=0, step=1, label="Seed")
-                        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
-                gr.Markdown(
-                    "<small>Output size auto-matches your photo's aspect ratio (~1MP). "
-                    "For two-person edits the LoRA works best ≤1.5MP.</small>"
-                )
-            with gr.Column(scale=1):
-                result = gr.Image(label="Result", format="png", height=420)
-                used_seed = gr.Number(label="Seed used", visible=True, interactive=False)
+                        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-        inputs = [
-            prompt, reference_image, use_reference,
-            lora_scale, steps, guidance, seed, randomize_seed,
-        ]
-        outputs = [result, used_seed]
+            with gr.Column():
+                result = gr.Image(label="Edited image", height=420)
+                used_seed = gr.Number(label="Seed used", interactive=False)
 
         gr.Examples(
             examples=[
-                ["place this on a wooden table with soft window light", "examples/style_ref_clean.png", True],
-                ["make the background a golden sunset", "examples/style_ref_clean.png", True],
+                ["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=[prompt, reference_image, use_reference],
-            outputs=outputs,
-            fn=generate,
+            inputs=[source_image, instruction],
+            outputs=[result, used_seed],
+            fn=edit,
             cache_examples=True,
             cache_mode="lazy",
         )
 
-        gr.on([generate_btn.click, prompt.submit], generate, inputs, outputs)
+    gr.on(
+        triggers=[run_button.click, instruction.submit],
+        fn=edit,
+        inputs=[source_image, instruction, 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)
+    demo.launch(theme=gr.themes.Citrus(), css=CSS, show_error=True)