"""Krea 2 (K2) for ai-toolkit. Krea 2 is a single-stream MMDiT text-to-image model: - text encoder: Qwen3-VL-4B-Instruct (a stack of 12 hidden-state layers is fed in; ``src/text_encoder.py``), - autoencoder: the Qwen-Image VAE (f8, 16 latent channels, the same VAE the ``qwen_image`` arch uses), - denoiser: ``SingleStreamDiT`` (``src/mmdit.py``), which fuses the text layers with a small ``TextFusionTransformer`` and runs the packed [text | image] sequence through ``SingleStreamBlock`` layers. Flow-matching convention matches ai-toolkit exactly (t=1 noise -> t=0 clean, target = noise - clean), so ``get_noise_prediction`` does no time flip / negation. """ import math import os from typing import TYPE_CHECKING, List, Optional import torch import torch.nn.functional as F from PIL import Image from torchvision.transforms.functional import to_tensor from safetensors.torch import load_file, save_file import huggingface_hub from huggingface_hub.errors import EntryNotFoundError from diffusers import AutoencoderKLQwenImage from transformers import ( AutoProcessor, AutoTokenizer, Qwen2TokenizerFast, Qwen3VLForConditionalGeneration, ) from optimum.quanto import freeze from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig from toolkit.lora_special import LoRASpecialNetwork from toolkit.models.base_model import BaseModel from toolkit.basic import flush from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.samplers.custom_flowmatch_sampler import ( CustomFlowMatchEulerDiscreteScheduler, ) from toolkit.accelerator import unwrap_model from toolkit.metadata import get_meta_for_safetensors from toolkit.util.quantize import quantize, get_qtype, quantize_model from toolkit.memory_management import MemoryManager from .src.mmdit import ( DoubleSharedModulation, SimpleModulation, SingleMMDiTConfig, SingleStreamDiT, ) from .src.text_encoder import encode_krea_prompt, SELECT_LAYERS from .src.pipeline import Krea2Pipeline, pad_text_features, predict_velocity if TYPE_CHECKING: from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO # The reference "single_mmdit_large_wide" architecture (oss_raw / oss_turbo share it). KREA2_MMDIT_CONFIG = dict( features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4, layers=28, patch=2, channels=16, txtheads=20, txtkvheads=20, txtlayers=12, ) # Krea 2's mu schedule is exponential time-shifting whose mu is linearly # interpolated in image-token count between (256-res -> 0.5) and (1280-res -> # 1.15) -- exactly what CustomFlowMatchEulerDiscreteScheduler's dynamic shifting # does, so we mirror those endpoints here for the training timestep distribution. # x1 = (256 // (8*2))**2 = 256 # x2 = (1280 // (8*2))**2 = 6400 scheduler_config = { "base_image_seq_len": 256, "max_image_seq_len": 6400, "base_shift": 0.5, "max_shift": 1.15, "num_train_timesteps": 1000, "shift": 1.0, "use_dynamic_shifting": True, "time_shift_type": "exponential", } # Defaults; both overridable via model.model_kwargs. QWEN3_VL_PATH = "Qwen/Qwen3-VL-4B-Instruct" QWEN_IMAGE_VAE_PATH = "Qwen/Qwen-Image" HF_TOKEN = os.getenv("HF_TOKEN", None) def patch_qwen_vl_patch_embed(model): """Qwen-VL's vision patch_embed is a Conv3d whose kernel == stride, i.e. a plain linear projection of each flattened patch. bf16 Conv3d has no fast cuDNN kernel and falls back to a slow, GPU-underutilizing path. Swap it for the equivalent F.linear (a GEMM). The weight is read lazily so this survives later .to(device)/dtype moves. Returns the number of patch_embed modules patched. (Same patch as the Qwen3VLCaptioner extension.)""" patched = 0 for module in model.modules(): proj = getattr(module, "proj", None) if isinstance(proj, torch.nn.Conv3d) and tuple(proj.kernel_size) == tuple( proj.stride ): def fast_forward(hidden_states, _proj=proj): w = _proj.weight.reshape(_proj.weight.shape[0], -1) x = hidden_states.view(-1, w.shape[1]).to(w.dtype) return F.linear(x, w, _proj.bias) module.forward = fast_forward patched += 1 return patched def _load_mmdit_state_dict(name_or_path: str, filename: Optional[str]) -> dict: """Load the MMDiT weights from a local safetensors file/dir or the HF hub. ``name_or_path`` may be: a ``.safetensors`` file, a directory containing one (``filename`` or the lone ``.safetensors`` in it), or a hub repo id (the file ``filename`` is downloaded, defaulting to ``model.safetensors``). """ if name_or_path.endswith(".safetensors") and os.path.isfile(name_or_path): return load_file(name_or_path) if os.path.isdir(name_or_path): if filename is not None: return load_file(os.path.join(name_or_path, filename)) candidates = [f for f in os.listdir(name_or_path) if f.endswith(".safetensors")] if len(candidates) == 1: return load_file(os.path.join(name_or_path, candidates[0])) raise FileNotFoundError( f"Could not pick an MMDiT checkpoint in {name_or_path}: found " f"{candidates}. Set model.model_kwargs.checkpoint_filename." ) # Treat as a hub repo id. When no filename is given, derive it from the repo # name's trailing segment (e.g. "krea/Krea-2-Raw" -> "raw.safetensors", # "krea/Krea-2-Turbo" -> "turbo.safetensors"). fname = filename or ( name_or_path.split("/")[-1].split("-")[-1].lower() + ".safetensors" ) try: path = huggingface_hub.hf_hub_download( repo_id=name_or_path, filename=fname, token=HF_TOKEN ) except EntryNotFoundError as e: raise FileNotFoundError( f"Could not find {fname!r} in hub repo {name_or_path!r}. Set " "model.model_kwargs.checkpoint_filename to the weight file name." ) from e return load_file(path) class Krea2Model(BaseModel): arch = "krea2" def __init__( self, device, model_config: ModelConfig, dtype="bf16", custom_pipeline=None, noise_scheduler=None, **kwargs, ): super().__init__( device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs ) self.is_flow_matching = True self.is_transformer = True self.target_lora_modules = ["SingleStreamDiT"] self.patch_size = KREA2_MMDIT_CONFIG["patch"] self.vae_scale_factor = 8 # Qwen-Image VAE is f8 # Safety cap on prompt token length (truncation only); embeds are stored # per-sample at natural length and padded to the batch max at the model call. self.max_text_length = int( self.model_config.model_kwargs.get("max_text_length", 512) ) # Qwen2TokenizerFast used to tokenize the assistant suffix (matches the # reference's separate processor pass). self.processor = None # Qwen3-VL AutoProcessor for encoding reference images into the prompt. self.vl_processor = None self.use_old_lokr_format = False # Optional reference-image (edit) conditioning, enabled with # model_kwargs.edit = true. Control images feed the model in two places: # through the Qwen3-VL encoder alongside the prompt (edit-plus style, so # the text embeddings see them) and as clean VAE latents appended to the # image sequence at t=0 (ComfyUI Kontext "index_timestep_zero"). Runs in # ComfyUI with the ComfyUI-Krea2-Ostris-Edit custom nodes. With edit off # (the default) all of it is skipped and this is the plain T2I model. self.is_edit = bool(self.model_config.model_kwargs.get("edit", False)) self.encode_control_in_text_embeddings = self.is_edit self.has_multiple_control_images = self.is_edit # Reference images keep their own aspect/size (not resized to the target). self.use_raw_control_images = self.is_edit # model_kwargs.kv_cache = true: train with an asymmetric attention mask # where the clean reference tokens attend only to each other (never to # text / noisy tokens). Their hidden states then depend only on the # refs + t=0 modulation, so at inference their per-layer K/V can be # computed once and reused across all denoising steps # (OminiControl2-style conditioning feature reuse). Off by default: # the base model was trained fully bidirectional, so a LoRA must be # trained with kv_cache enabled for kv-cached inference (the ComfyUI # node / hub pipeline kv_cache toggles) to work properly. self.kv_cache = bool(self.model_config.model_kwargs.get("kv_cache", False)) @staticmethod def get_train_scheduler(): return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) def get_bucket_divisibility(self): # 8 for the VAE downsample, 2 for the patch size. return self.vae_scale_factor * self.patch_size # ------------------------------------------------------------------ # Loading # ------------------------------------------------------------------ def _load_transformer(self): dtype = self.torch_dtype self.print_and_status_update("Loading transformer (SingleStreamDiT)") mmdit_kwargs = dict(KREA2_MMDIT_CONFIG) mmdit_kwargs.update(self.model_config.model_kwargs.get("mmdit_config", {})) config = SingleMMDiTConfig(**mmdit_kwargs) # Build on meta, then materialize straight from the checkpoint. with torch.device("meta"): transformer = SingleStreamDiT(config) self.print_and_status_update(" - fetching transformer weights") state_dict = _load_mmdit_state_dict( self.model_config.name_or_path, self.model_config.model_kwargs.get("checkpoint_filename", None), ) state_dict = { k: (v.to(dtype) if v.is_floating_point() else v) for k, v in state_dict.items() } self.print_and_status_update(" - loading transformer state dict") transformer.load_state_dict(state_dict, strict=True, assign=True) del state_dict flush() return transformer def _load_text_encoder(self): dtype = self.torch_dtype te_path = self.model_config.model_kwargs.get("text_encoder_path", QWEN3_VL_PATH) self.print_and_status_update(f"Loading Qwen3-VL text encoder from {te_path}") tokenizer = AutoTokenizer.from_pretrained( te_path, max_length=self.max_text_length, token=HF_TOKEN ) processor = Qwen2TokenizerFast.from_pretrained( te_path, max_length=self.max_text_length, token=HF_TOKEN ) text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( te_path, torch_dtype=dtype, token=HF_TOKEN ) vl_processor = None if self.is_edit: # Edit mode: reference images are encoded into the text embeddings, # so the vision tower stays. Swap its Conv3d patch_embed for an # equivalent GEMM (bf16 Conv3d has no fast cuDNN kernel). vl_processor = AutoProcessor.from_pretrained(te_path, token=HF_TOKEN) patch_qwen_vl_patch_embed(text_encoder) else: # We only ever encode text, so the vision tower is dead weight -- drop it to # free VRAM and skip loading its (bf16-slow) Conv3d patch_embed onto the GPU. if getattr(text_encoder.model, "visual", None) is not None: text_encoder.model.visual = None text_encoder.eval() text_encoder.requires_grad_(False) flush() return tokenizer, processor, vl_processor, text_encoder def _load_vae(self): vae_path = self.model_config.model_kwargs.get("vae_path", QWEN_IMAGE_VAE_PATH) self.print_and_status_update(f"Loading Qwen-Image VAE from {vae_path}") vae = AutoencoderKLQwenImage.from_pretrained( vae_path, subfolder="vae", torch_dtype=self.vae_torch_dtype, token=HF_TOKEN ) vae.eval() vae.requires_grad_(False) return vae def load_training_adapter(self, transformer: SingleStreamDiT): self.print_and_status_update("Loading assistant LoRA") lora_path = self.model_config.assistant_lora_path if not os.path.exists(lora_path): # assume it is a hub path lora_splits = lora_path.split("/") if len(lora_splits) != 3: raise ValueError( f"Assistant LoRA path {lora_path} is not a valid local path or hub path." ) repo_id = "/".join(lora_splits[:2]) filename = lora_splits[2] try: lora_path = huggingface_hub.hf_hub_download( repo_id=repo_id, filename=filename, token=HF_TOKEN, ) # upgrade path to the local download self.model_config.assistant_lora_path = lora_path except Exception as e: raise ValueError( f"Failed to download assistant LoRA from {lora_path}: {e}" ) # load the adapter and merge it in. We will inference with a -1.0 multiplier so the adapter effects only work during training. lora_state_dict = load_file(lora_path) # detect the LoRA rank from the first down-projection weight. dim_key = next(k for k in lora_state_dict if k.endswith("lora_A.weight")) dim = int(lora_state_dict[dim_key].shape[0]) new_sd = {} for key, value in lora_state_dict.items(): new_key = key.replace("diffusion_model.", "transformer.") new_sd[new_key] = value lora_state_dict = new_sd network_config = { "type": "lora", "linear": dim, "linear_alpha": dim, "transformer_only": True, } network_config = NetworkConfig(**network_config) LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer" network = LoRASpecialNetwork( text_encoder=None, unet=transformer, lora_dim=network_config.linear, multiplier=1.0, alpha=network_config.linear_alpha, train_unet=True, train_text_encoder=False, network_config=network_config, network_type=network_config.type, transformer_only=network_config.transformer_only, is_transformer=True, target_lin_modules=self.target_lora_modules, is_assistant_adapter=True, is_ara=True, ) network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True) self.print_and_status_update("Merging in assistant LoRA") network.force_to(self.device_torch, dtype=self.torch_dtype) network._update_torch_multiplier() network.load_weights(lora_state_dict) network.merge_in(merge_weight=1.0) # mark it as not merged so inference ignores it. network.is_merged_in = False # add the assistant so sampler will activate it while sampling self.assistant_lora: LoRASpecialNetwork = network # deactivate lora during training self.assistant_lora.multiplier = -1.0 self.assistant_lora.is_active = False # tell the model to invert assistant on inference since we want remove lora effects self.invert_assistant_lora = True def load_model(self): dtype = self.torch_dtype self.print_and_status_update("Loading Krea 2 model") transformer = self._load_transformer() # load assistant lora if specified if self.model_config.assistant_lora_path is not None: self.load_training_adapter(transformer) # set qtype to be float8 if it is qfloat8 if self.model_config.qtype == "qfloat8": self.model_config.qtype = "float8" if self.model_config.quantize: self.print_and_status_update("Quantizing transformer") quantize_model(self, transformer) flush() if ( self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0 ): MemoryManager.attach( transformer, self.device_torch, offload_percent=self.model_config.layer_offloading_transformer_percent, ignore_modules=[ module for module in transformer.modules() if isinstance(module, (SimpleModulation, DoubleSharedModulation)) ], ) if self.model_config.low_vram: self.print_and_status_update("Moving transformer to CPU") transformer.to("cpu") else: transformer.to(self.device_torch, dtype=dtype) flush() tokenizer, processor, vl_processor, text_encoder = self._load_text_encoder() if self.model_config.quantize_te: self.print_and_status_update("Quantizing text encoder") text_encoder.to(self.device_torch) quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) freeze(text_encoder) flush() if ( self.model_config.layer_offloading and self.model_config.layer_offloading_text_encoder_percent > 0 ): MemoryManager.attach( text_encoder, self.device_torch, offload_percent=self.model_config.layer_offloading_text_encoder_percent, ) if self.model_config.low_vram: self.print_and_status_update("Moving text encoder to CPU") text_encoder.to("cpu") else: text_encoder.to(self.device_torch) flush() vae = self._load_vae() vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype) self.noise_scheduler = Krea2Model.get_train_scheduler() self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer self.processor = processor self.vl_processor = vl_processor self.model = transformer self.pipeline = Krea2Pipeline(self) self.print_and_status_update("Model Loaded") # ------------------------------------------------------------------ # Generation (training previews) # ------------------------------------------------------------------ def get_generation_pipeline(self): return Krea2Pipeline(self) def generate_single_image( self, pipeline: Krea2Pipeline, gen_config: GenerateImageConfig, conditional_embeds: AdvancedPromptEmbeds, unconditional_embeds: AdvancedPromptEmbeds, generator: torch.Generator, extra: dict, ): if self.model.device == torch.device("cpu"): self.model.to(self.device_torch) sc = self.get_bucket_divisibility() gen_config.width = int(gen_config.width // sc * sc) gen_config.height = int(gen_config.height // sc * sc) # Reference image(s) -> clean VAE latents for the t=0 sequence tokens. # The Qwen3-VL side already saw them (baked into the prompt embeds). # ctrl_img_1 mirrors ctrl_img when unset, so use one or the other. ctrl_paths = [] if self.is_edit: if gen_config.ctrl_img is not None: ctrl_paths.append(gen_config.ctrl_img) elif gen_config.ctrl_img_1 is not None: ctrl_paths.append(gen_config.ctrl_img_1) if gen_config.ctrl_img_2 is not None: ctrl_paths.append(gen_config.ctrl_img_2) if gen_config.ctrl_img_3 is not None: ctrl_paths.append(gen_config.ctrl_img_3) ref_latents = None if ctrl_paths: ctrl_tensors = [ to_tensor(Image.open(path).convert("RGB")) for path in ctrl_paths ] target_pixels = gen_config.width * gen_config.height # one batch item (preview batch size is 1) -> List[List[(16, h, w)]] ref_latents = [ self._encode_ref_latents(ctrl_tensors, target_pixels=target_pixels) ] # CFG is 0 normalized for this model guidance = max(0.0, gen_config.guidance_scale - 1.0) img = pipeline( conditional_embeds=conditional_embeds, unconditional_embeds=unconditional_embeds, height=gen_config.height, width=gen_config.width, num_inference_steps=gen_config.num_inference_steps, guidance_scale=guidance, latents=gen_config.latents, generator=generator, ref_latents=ref_latents, )[0] return img # ------------------------------------------------------------------ # Reference-image helpers # ------------------------------------------------------------------ def _ref_target_pixels(self, target_pixels: Optional[int]) -> int: """Pixel budget each reference image is resized to fit within. - default: ``control_image_max_pixels`` model_kwarg (1 MP) -- a hard cap so raw, full-size control images don't blow up the token count / VRAM. - ``match_target_res`` model_kwarg: use the target generation area instead. """ max_pixels = int( self.model_config.model_kwargs.get("control_image_max_pixels", 1024 * 1024) ) if ( self.model_config.model_kwargs.get("match_target_res", False) and target_pixels ): return int(target_pixels) return max_pixels def _encode_ref_latents( self, control_tensors, target_pixels: Optional[int] = None ) -> List[torch.Tensor]: """Encode ``[0, 1]`` reference image tensors to VAE latents. Returns a list of ``(16, h, w)`` latents (one per reference image). Each control image is resized so its area fits within the pixel budget (see ``_ref_target_pixels``) -- preserving aspect ratio -- then snapped so the latent grid is divisible by the patch size. ``control_tensors`` is a list of ``(C, H, W)`` or ``(1, C, H, W)`` tensors in ``[0, 1]``. """ sc = self.get_bucket_divisibility() # 16: VAE(8) * patch(2) budget = self._ref_target_pixels(target_pixels) match = self.model_config.model_kwargs.get("match_target_res", False) latents = [] for img in control_tensors: if img.dim() == 3: img = img.unsqueeze(0) img = img.to(self.device_torch, dtype=self.torch_dtype) h, w = img.shape[2], img.shape[3] # match_target_res: scale area *to* the budget; otherwise only scale # *down* when the image is larger than the budget. area = h * w if match or area > budget: ratio = h / w new_h = math.sqrt(budget * ratio) new_w = new_h / ratio else: new_h, new_w = float(h), float(w) # snap to a multiple of the bucket divisibility so the VAE latent grid # is patchifiable (the transformer rearranges 2x2 latent patches). new_h = max(sc, int(round(new_h / sc)) * sc) new_w = max(sc, int(round(new_w / sc)) * sc) if (new_h, new_w) != (h, w): img = F.interpolate(img, size=(new_h, new_w), mode="bilinear") # encode_images expects [-1, 1]; control tensors arrive in [0, 1]. latent = self.encode_images( img * 2 - 1, device=self.device_torch, dtype=self.torch_dtype ) latents.append(latent[0]) # drop batch dim -> (16, h, w) return latents def _batch_ref_latents_from_batch( self, batch: "DataLoaderBatchDTO", batch_size: int, target_pixels: Optional[int] = None, ) -> Optional[List[List[torch.Tensor]]]: """Build predict_velocity's ``ref_latents`` from a train batch.""" control_list = batch.control_tensor_list if control_list is None and batch.control_tensor is not None: control_list = [batch.control_tensor[b : b + 1] for b in range(batch_size)] if control_list is None: return None if len(control_list) != batch_size: raise ValueError("Control tensor list length does not match batch size") return [ self._encode_ref_latents(controls, target_pixels=target_pixels) for controls in control_list ] # ------------------------------------------------------------------ # Training hooks # ------------------------------------------------------------------ def get_noise_prediction( self, latent_model_input: torch.Tensor, # (B, 16, h, w) timestep: torch.Tensor, # 0..1000 scale text_embeddings: AdvancedPromptEmbeds, batch: "DataLoaderBatchDTO" = None, **kwargs, ): if self.model.device == torch.device("cpu"): self.model.to(self.device_torch) # Clean reference latents from the batch's control images (if any); they # ride along in the sequence at t=0 and are never noised. ref_latents = None if batch is not None and self.is_edit: with torch.no_grad(): _, _, lh, lw = latent_model_input.shape target_pixels = (lh * self.vae_scale_factor) * ( lw * self.vae_scale_factor ) ref_latents = self._batch_ref_latents_from_batch( batch, latent_model_input.shape[0], target_pixels=target_pixels ) # toolkit timestep (0..1000, 1000 = pure noise) -> Krea flow time t in # [0, 1] with t=1 = pure noise. Same convention -> straight divide. t = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0 if t.dim() == 0: t = t.unsqueeze(0) if t.shape[0] != latent_model_input.shape[0]: t = t.expand(latent_model_input.shape[0]) context, text_mask = pad_text_features( text_embeddings.text_embeds, self.device_torch, self.torch_dtype ) pred = predict_velocity( self.transformer, latent_model_input.to(self.device_torch, self.torch_dtype), t, context, text_mask, ref_latents=ref_latents, isolate_refs=self.kv_cache, ) return pred def _prep_vlm_images(self, ctrl: List[torch.Tensor]) -> List[torch.Tensor]: """Resize reference images for the Qwen3-VL pass. Downscaled (aspect-preserved, never upscaled) to fit ``vlm_max_pixels`` total area (384^2 by default, the boogu_image_edit / ComfyUI TextEncodeQwenImageEditPlus budget) -- the MLLM only needs a coarse understanding of the reference; high-res detail flows through the VAE ref latents. """ target = int(self.model_config.model_kwargs.get("vlm_max_pixels", 384 * 384)) images = [] for img in ctrl: if img.dim() == 4: img = img[0] img = img.to(self.device_torch) h, w = img.shape[1], img.shape[2] scale = min(1.0, math.sqrt(target / (h * w))) nh, nw = max(round(h * scale), 28), max(round(w * scale), 28) if (nh, nw) != (h, w): img = ( F.interpolate( img.unsqueeze(0).float(), size=(nh, nw), mode="bicubic", antialias=True, ) .squeeze(0) .clamp(0, 1) ) images.append(img.float()) return images def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds: if isinstance(prompt, str): prompt = [prompt] if self.text_encoder.device == torch.device("cpu"): self.text_encoder.to(self.device_torch) # Normalize control images to a per-prompt list (List[List[Tensor]]). # They arrive as a (B, C, H, W) batch tensor (control_tensor), a list of # per-sample lists (control_tensor_list), or a flat list of (1, C, H, W) # tensors for a single prompt (sampling / blank-embed caching). if control_images is not None: if isinstance(control_images, torch.Tensor): control_images = [ [control_images[i]] for i in range(control_images.shape[0]) ] elif len(control_images) > 0 and not isinstance(control_images[0], list): control_images = [control_images] if len(control_images) == 1 and len(prompt) > 1: control_images = control_images * len(prompt) if len(control_images) != len(prompt): raise ValueError( "Number of prompts must match number of control image sets" ) else: control_images = [None] * len(prompt) # Encode each prompt at its natural length and store one (L, 12*2560) # tensor per batch item. The (L, 12, 2560) stack is flattened to 2D so the # toolkit's batching reads the list length (not the seq length) as the # batch size; predict_velocity restores the layer axis. Padding to the # batch max is deferred to the model call so caches stay small and any # prompts can share a batch. features_list = [] for p, ctrl in zip(prompt, control_images): images = self._prep_vlm_images(ctrl) if ctrl is not None else None features = encode_krea_prompt( self.text_encoder, self.tokenizer, self.processor, p, max_length=self.max_text_length, select_layers=SELECT_LAYERS, images=images, vl_processor=self.vl_processor, dtype=self.torch_dtype, ) # (L, n, d) -> (L, n*d) features = features.reshape(features.shape[0], -1) features_list.append(features.to(self.torch_dtype)) return AdvancedPromptEmbeds(text_embeds=features_list) def get_loss_target(self, *args, **kwargs): # Flow-matching velocity target: noise - clean. noise = kwargs.get("noise") batch = kwargs.get("batch") return (noise - batch.latents).detach() def get_model_has_grad(self): return False def get_te_has_grad(self): return False # ------------------------------------------------------------------ # VAE (Qwen-Image AutoencoderKLQwenImage -- same handling as qwen_image arch) # ------------------------------------------------------------------ def encode_images(self, image_list: List[torch.Tensor], device=None, dtype=None): if device is None: device = self.vae_device_torch if dtype is None: dtype = self.vae_torch_dtype if self.vae.device == torch.device("cpu"): self.vae.to(device) self.vae.eval() self.vae.requires_grad_(False) image_list = [image.to(device, dtype=dtype) for image in image_list] images = torch.stack(image_list).to(device, dtype=dtype) # AutoencoderKLQwenImage is a video VAE: add a frame dim. images = images.unsqueeze(2) latents = self.vae.encode(images).latent_dist.sample() latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(latents.device, latents.dtype) ) latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view( 1, self.vae.config.z_dim, 1, 1, 1 ).to(latents.device, latents.dtype) latents = (latents - latents_mean) * latents_std latents = latents.squeeze(2) # drop frame dim return latents.to(device, dtype=dtype) def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): if device is None: device = self.vae_device_torch if dtype is None: dtype = self.vae_torch_dtype if self.vae.device == torch.device("cpu"): self.vae.to(device) latents = latents.to(device, dtype=dtype) latents = latents.unsqueeze(2) # add frame dim latents_mean = ( torch.tensor(self.vae.config.latents_mean) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(latents.device, latents.dtype) ) latents_std = ( torch.tensor(self.vae.config.latents_std) .view(1, self.vae.config.z_dim, 1, 1, 1) .to(latents.device, latents.dtype) ) latents = latents * latents_std + latents_mean # Full-resolution decode spikes VRAM; tile it when low on VRAM (decode # only -- encode stays untiled). tiled = self.model_config.low_vram if tiled: self.vae.enable_tiling() try: images = self.vae.decode(latents).sample finally: if tiled: self.vae.disable_tiling() images = images.squeeze(2) # drop frame dim return images.to(device, dtype=dtype) # ------------------------------------------------------------------ # Saving / bookkeeping # ------------------------------------------------------------------ def save_model(self, output_path, meta, save_dtype): from toolkit.util.quantize import dequantize_if_quantized if not output_path.endswith(".safetensors"): output_path = output_path + ".safetensors" transformer: SingleStreamDiT = unwrap_model(self.model) state_dict = transformer.state_dict() save_dict = {} for k, v in state_dict.items(): # dequantize any quantized (e.g. quanto/torchao) weights so we save plain full precision tensors save_dict[k] = ( dequantize_if_quantized(v).clone().to("cpu", dtype=save_dtype) ) meta = get_meta_for_safetensors(meta, name="krea2") save_file(save_dict, output_path, metadata=meta) def get_base_model_version(self): return "krea2" def get_transformer_block_names(self) -> Optional[List[str]]: return ["blocks"] def convert_lora_weights_before_save(self, state_dict): return { k.replace("transformer.", "diffusion_model."): v for k, v in state_dict.items() } def convert_lora_weights_before_load(self, state_dict): return { k.replace("diffusion_model.", "transformer."): v for k, v in state_dict.items() }