Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/man_beach.jpg filter=lfs diff=lfs merge=lfs -text
+examples/woman.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,49 @@
 ---
-title: Krea2 Identity Edit
-emoji: 📚
-colorFrom: red
-colorTo: yellow
+title: Krea 2 Identity Edit
+emoji: 🧬
+colorFrom: blue
+colorTo: gray
 sdk: gradio
 sdk_version: 6.20.0
-python_version: '3.13'
 app_file: app.py
 pinned: false
+hardware: zero-a10g
+python_version: "3.10"
+startup_duration_timeout: 45m
+short_description: Identity-preserving instruction image editing on Krea 2
+models:
+  - krea/Krea-2-Turbo
+  - conradlocke/krea2-identity-edit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🧬 Krea 2 Identity Edit
+
+Instruction-based, **identity-preserving** image editing built 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).
+
+Give it an image and a plain-language instruction; it edits while preserving what you
+didn't ask to change — including the person's likeness. Person re-staging, local edits
+(recolor / add / remove / replace), replace-with-reference, and full-image restyles.
+
+## How it works
+
+Krea 2 is a text-to-image MMDiT; this LoRA turns it into an instruction editor via a
+**dual-conditioning** recipe that stock `Krea2Pipeline` does not provide. This Space
+reproduces the two custom pieces from the reference
+[ComfyUI-Krea2Edit](https://github.com/lbouaraba/comfyui-krea2edit) node pack in plain
+diffusers / PyTorch:
+
+- **Grounded encode** — the instruction is encoded *together with the source image*
+  through the Qwen3-VL text encoder (image tokens inserted via the grounded chat
+  template), tapping the 12 selected decoder layers.
+- **Source patch** — the VAE-encoded source latent is prepended to the transformer
+  sequence as clean tokens (RoPE frame index: source = 1, target = 0), so appearance is
+  carried in-context and only the target tokens are denoised.
+
+Default recipe: Turbo, 8 steps, guidance off (Krea convention). `grounding_px` trades
+edit adherence (lower) against identity/likeness (higher); 768 is balanced, 1024+ for
+people. For removals / large deletions, raise steps and guidance.
+
+*An unofficial community fine-tune of Krea 2 Raw — not affiliated with or endorsed by
+Krea.ai, Inc. Weights are distributed under the Krea 2 Community License.*

app.py

@@ -0,0 +1,444 @@
+"""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
+import random
+
+import numpy as np
+import spaces
+import torch
+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.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, 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 = 2.0            # card: generate at <= 2MP
+
+# 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)
+
+# Load the identity-edit LoRA onto the transformer (Krea 2 LoRAs load through the
+# transformer's adapter API, per the base-model reference and the LoRA cards).
+pipe.transformer.load_lora_adapter(LORA_REPO, weight_name=LORA_WEIGHT)
+pipe.transformer.set_adapters("default", weights=DEFAULT_LORA_SCALE)
+
+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: (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)
+
+    outputs = pipe.text_encoder(
+        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,
+    )
+    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
+    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)
+
+    # Attention mask: text uses its key-padding mask, all image tokens (src+tgt) valid.
+    attention_mask = None
+    if prompt_mask is not None:
+        img_ones = prompt_mask.new_ones((combined_img.shape[0], combined_img.shape[1]))
+        attention_mask = torch.cat([prompt_mask, img_ones], dim=1)[:, None, None, :]
+
+    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(duration=90)
+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 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")
+
+    pipe.transformer.set_adapters("default", weights=float(lora_scale))
+
+    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
+        )
+        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_recip = (1.0 / _LATENTS_STD).to(latents.device, latents.dtype)
+    latents = latents * std_recip + 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).
+
+Upload a photo of a person (or any scene), type a plain-language instruction, and the
+model edits it while preserving what you didn't ask to change — **including the face**.
+Try *"put this person at a busy night market"*, *"change the jacket to red leather"*, or
+*"make it a vintage film photo"*.
+
+This demo reproduces the dual-conditioning recipe (in-context VAE source tokens +
+image-grounded Qwen3-VL encoding) from the reference
+[ComfyUI-Krea2Edit](https://github.com/lbouaraba/comfyui-krea2edit) node pack.
+"""
+
+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")
+
+                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",
+                        minimum=0.0, maximum=1.5, step=0.05, value=DEFAULT_LORA_SCALE,
+                    )
+                    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, 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)

requirements.txt

@@ -0,0 +1,9 @@
+git+https://github.com/huggingface/diffusers
+transformers
+accelerate
+peft
+safetensors
+sentencepiece
+torchvision
+Pillow
+numpy