"""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_1.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 = 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() 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 _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 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 ) 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 = """ #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). """ 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 (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, 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)