"""Klein text encoder for Flux2 models in LightDiffusion-Next. This module provides the Klein (Qwen3-4B based) text encoder used by Flux2 Klein models, including: - KleinTokenizer: Qwen3-based tokenizer with special formatting - Qwen3Model: Transformer-based language model for text encoding Adapted from ComfyUI's Klein implementation. """ import logging import math from typing import Optional, List, Dict, Any import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from src.cond import cast as ops_module from src.Device import Device logger = logging.getLogger(__name__) def get_ops(): """Get the operations module for weight initialization.""" return ops_module.disable_weight_init class QwenRMSNorm(nn.Module): """RMS Normalization for Qwen3.""" def __init__(self, dim: int, eps: float = 1e-6, dtype=None, device=None, operations=None): super().__init__() if operations is None: operations = get_ops() self.eps = eps self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=device)) def forward(self, x: torch.Tensor) -> torch.Tensor: # RMS normalization - compute in float32 for precision, cast back to input dtype input_dtype = x.dtype x_float = x.float() rms = torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps) return (x_float * rms * self.weight.float()).to(input_dtype) class QwenRotaryEmbedding(nn.Module): """Rotary position embeddings for Qwen3.""" def __init__(self, dim: int, max_position_embeddings: int = 32768, base: float = 1000000.0): super().__init__() self.dim = dim self.max_seq_len = max_position_embeddings self.base = base inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, x: torch.Tensor, seq_len: int = None): if seq_len is None: seq_len = x.shape[1] t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype) freqs = torch.outer(t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos(), emb.sin() def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotate half for RoPE.""" x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): """Apply rotary position embeddings to query and key.""" cos = cos.unsqueeze(0).unsqueeze(0) # [1, 1, seq, dim] sin = sin.unsqueeze(0).unsqueeze(0) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class QwenAttention(nn.Module): """Multi-head attention for Qwen3 with Grouped Query Attention.""" def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: int = None, head_dim: int = 128, dtype=None, device=None, operations=None, ): super().__init__() if operations is None: operations = get_ops() self.hidden_size = hidden_size self.num_heads = num_heads self.num_kv_heads = num_kv_heads or num_heads self.head_dim = head_dim # Qwen3 uses separate projections with different output sizes self.q_proj = operations.Linear(hidden_size, num_heads * head_dim, bias=False, dtype=dtype, device=device) self.k_proj = operations.Linear(hidden_size, self.num_kv_heads * head_dim, bias=False, dtype=dtype, device=device) self.v_proj = operations.Linear(hidden_size, self.num_kv_heads * head_dim, bias=False, dtype=dtype, device=device) self.o_proj = operations.Linear(num_heads * head_dim, hidden_size, bias=False, dtype=dtype, device=device) # Normalize Q and K self.q_norm = QwenRMSNorm(head_dim, dtype=dtype, device=device, operations=operations) self.k_norm = QwenRMSNorm(head_dim, dtype=dtype, device=device, operations=operations) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[tuple] = None, ) -> torch.Tensor: batch_size, seq_len, _ = hidden_states.shape q = self.q_proj(hidden_states) k = self.k_proj(hidden_states) v = self.v_proj(hidden_states) # Reshape for multi-head attention q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) # Apply QK normalization q = self.q_norm(q) k = self.k_norm(k) # Apply rotary embeddings if position_embeddings is not None: cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) # Grouped query attention - repeat K,V for each group if self.num_kv_heads != self.num_heads: n_rep = self.num_heads // self.num_kv_heads k = k.repeat_interleave(n_rep, dim=1) v = v.repeat_interleave(n_rep, dim=1) # Ensure all tensors have same dtype for SDPA attn_dtype = q.dtype k = k.to(attn_dtype) v = v.to(attn_dtype) # Scaled dot-product attention with causal masking # Use is_causal=True for efficiency, or attn_mask for custom masks if attention_mask is None: # Pure causal masking attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: # Custom mask (includes causal + padding) - ensure mask dtype matches attention_mask = attention_mask.to(attn_dtype) attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask) # Reshape back and ensure output dtype matches input for o_proj attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) attn_output = attn_output.to(hidden_states.dtype) # Match input dtype for o_proj return self.o_proj(attn_output) class QwenMLP(nn.Module): """MLP (Gate-Up-Down) for Qwen3.""" def __init__( self, hidden_size: int, intermediate_size: int, dtype=None, device=None, operations=None, ): super().__init__() if operations is None: operations = get_ops() self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device) self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class QwenDecoderLayer(nn.Module): """Single transformer decoder layer for Qwen3.""" def __init__( self, hidden_size: int, num_heads: int, intermediate_size: int, num_kv_heads: int = None, head_dim: int = 128, dtype=None, device=None, operations=None, ): super().__init__() if operations is None: operations = get_ops() self.self_attn = QwenAttention( hidden_size=hidden_size, num_heads=num_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, dtype=dtype, device=device, operations=operations, ) self.mlp = QwenMLP( hidden_size=hidden_size, intermediate_size=intermediate_size, dtype=dtype, device=device, operations=operations, ) self.input_layernorm = QwenRMSNorm(hidden_size, dtype=dtype, device=device, operations=operations) self.post_attention_layernorm = QwenRMSNorm(hidden_size, dtype=dtype, device=device, operations=operations) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[tuple] = None, ) -> torch.Tensor: # Self attention residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, attention_mask, position_embeddings) hidden_states = residual + hidden_states # MLP residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class Qwen3_4BModel(nn.Module): """Qwen3 4B model for Klein text encoding. This is a decoder-only transformer used as a text encoder for the Flux2 Klein model. """ def __init__( self, vocab_size: int = 151936, hidden_size: int = 2560, intermediate_size: int = 9728, # Matches checkpoint num_hidden_layers: int = 36, num_attention_heads: int = 32, # Matches checkpoint (4096/128) num_key_value_heads: int = 8, # Matches checkpoint (1024/128) head_dim: int = 128, max_position_embeddings: int = 32768, layer_indices: tuple = (9, 18, 27), # Layers to extract embeddings from dtype=None, device=None, operations=None, ): super().__init__() if operations is None: operations = get_ops() self.vocab_size = vocab_size self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.layer_indices = layer_indices # Token embeddings self.embed_tokens = operations.Embedding(vocab_size, hidden_size, dtype=dtype, device=device) # Rotary embeddings self.rotary_emb = QwenRotaryEmbedding( head_dim, max_position_embeddings=max_position_embeddings, ) # Transformer layers self.layers = nn.ModuleList([ QwenDecoderLayer( hidden_size=hidden_size, num_heads=num_attention_heads, intermediate_size=intermediate_size, num_kv_heads=num_key_value_heads, head_dim=head_dim, dtype=dtype, device=device, operations=operations, ) for _ in range(num_hidden_layers) ]) # Final norm self.norm = QwenRMSNorm(hidden_size, dtype=dtype, device=device, operations=operations) def forward( self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> dict: """Forward pass returning hidden states from specified layers. Args: input_ids: Token IDs [batch, seq_len] attention_mask: Optional attention mask Returns: Dict with 'hidden_states' from specified layers (concatenated) """ batch_size, seq_len = input_ids.shape # Embed tokens - ensure output is at least bfloat16 for subsequent math hidden_states = self.embed_tokens(input_ids).to(torch.bfloat16) # Get rotary embeddings cos, sin = self.rotary_emb(hidden_states, seq_len) position_embeddings = (cos, sin) # Prepare attention mask # If we have a padding mask, create a combined causal + padding mask # Otherwise, pass None and let the attention layer use is_causal=True final_mask = None if attention_mask is not None: # Create mask matching ComfyUI's approach: # 1. Convert padding mask from [B, L] to [B, 1, L, L] with expansion # 2. Set padded positions (where mask=0) to -inf # 3. Add causal mask # Reshape and expand: [B, L] -> [B, 1, L, L] mask = 1.0 - attention_mask.to(hidden_states.dtype) # 0 = valid, 1 = padding mask = mask.reshape(mask.shape[0], 1, -1, mask.shape[-1]) # [B, 1, 1, L] mask = mask.expand(mask.shape[0], 1, seq_len, seq_len) # [B, 1, L, L] mask = mask.masked_fill(mask.to(torch.bool), float("-inf")) # Create causal mask [L, L] causal_mask = torch.empty(seq_len, seq_len, dtype=hidden_states.dtype, device=input_ids.device).fill_(float("-inf")).triu_(1) # Combine final_mask = mask + causal_mask # Collect outputs from specified layers # NOTE: ComfyUI captures the INPUT to layers (before the layer runs), # so we capture before applying each layer layer_outputs = [] for i, layer in enumerate(self.layers): hidden_states = layer(hidden_states, final_mask, position_embeddings) # Capture AFTER the layer (output of layer i) if i in self.layer_indices: layer_outputs.append(hidden_states.clone()) # Apply final norm hidden_states = self.norm(hidden_states) # Concatenate layer outputs matching ComfyUI's interleaving pattern # This is critical for Flux2/Klein cross-attention if layer_outputs: # layer_outputs is a list of [B, L, D] tensors # stack: (B, 3, L, D) stacked = torch.stack(layer_outputs, dim=1) # permute: (B, L, 3, D) - interleave the 3 layers at each sequence position permuted = stacked.permute(0, 2, 1, 3) # reshape: (B, L, 3*D) concatenated = permuted.reshape(batch_size, seq_len, -1) else: concatenated = hidden_states return { "last_hidden_state": hidden_states, "hidden_states": concatenated, "pooled_output": None, # Match ComfyUI: No pooling for Qwen -> Flux2 uses zeros } class KleinTokenizer: """Tokenizer for Klein (Qwen3-based) text encoder. Uses Qwen2Tokenizer from Hugging Face transformers with Klein-specific formatting template. """ # Klein template for prompt formatting TEMPLATE = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" def __init__( self, tokenizer_path: str = None, max_length: int = 99999999, # ComfyUI uses essentially unlimited min_length: int = 512, # ComfyUI uses min_length=512 for Klein padding: str = "do_not_pad", # ComfyUI uses pad_to_max_length=False ): self.max_length = max_length self.min_length = min_length self.padding = padding # Klein special tokens self.pad_token_id = 151643 # <|endoftext|> self.bos_token_id = 151644 # <|im_start|> self.eos_token_id = 151645 # <|im_end|> # Load the real tokenizer if tokenizer_path is None: # Default path relative to include folder import os # Try multiple locations possible_paths = [ os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))), "include", "text_encoder", "qwen25_tokenizer"), os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer"), ] for path in possible_paths: if os.path.exists(path): tokenizer_path = path break else: tokenizer_path = possible_paths[0] # Use first as default try: from transformers import Qwen2Tokenizer self._tokenizer = Qwen2Tokenizer.from_pretrained(tokenizer_path) # Use right padding for content-first alignment, matching ComfyUI default self._tokenizer.padding_side = "right" logger.info(f"Loaded Qwen2Tokenizer from {tokenizer_path}") except Exception as e: logger.error(f"Failed to load tokenizer: {e}") raise RuntimeError(f"Could not load Klein tokenizer from {tokenizer_path}") from e def apply_template(self, text: str) -> str: """Apply Klein's prompt template to input text.""" return self.TEMPLATE.format(text) def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict: """Tokenize text with Klein template formatting. Args: text: Input text to tokenize return_word_ids: Whether to return word IDs Returns: Dict with 'input_ids' and 'attention_mask' """ # Apply template formatted_text = self.apply_template(text) # Tokenize with the real tokenizer - pad to min_length (512) # Matches ComfyUI Qwen3Tokenizer behavior encoded = self._tokenizer( formatted_text, padding="max_length", max_length=self.min_length, truncation=True, return_tensors="pt", ) result = { "input_ids": encoded["input_ids"], "attention_mask": encoded["attention_mask"], } if return_word_ids: # Word IDs from the tokenizer's encoding word_ids = encoded.word_ids() if hasattr(encoded, 'word_ids') else list(range(encoded["input_ids"].shape[1])) result["word_ids"] = word_ids return result def state_dict(self) -> dict: """Return tokenizer state for serialization.""" return { "max_length": self.max_length, "min_length": self.min_length, "padding": self.padding, } class KleinCLIP: """Klein text encoder wrapper compatible with CLIP interface. This provides the same interface as other CLIP models while using the Qwen3-based Klein encoder internally. VRAM Optimization: Model stays on CPU until encoding, then moves to GPU and back to CPU. This follows ComfyUI's lazy loading approach. """ def __init__( self, tokenizer: KleinTokenizer = None, model: Qwen3_4BModel = None, dtype=None, device=None, offload_device=None, ): self.tokenizer = tokenizer or KleinTokenizer() self.dtype = dtype self.device = device # Device to use for encoding (GPU) self.offload_device = offload_device or torch.device("cpu") # Device when idle (CPU) if model is None: self.model = Qwen3_4BModel(dtype=dtype, device=self.offload_device) else: self.model = model self.clip_options = {} def reset_clip_options(self): """Reset clip options to defaults.""" self.clip_options = {} def set_clip_options(self, options: dict): """Set clip options (for API compatibility).""" self.clip_options.update(options) def encode_token_weights(self, tokens: dict) -> tuple: """Encode token weights returning (embeddings, pooled, extra). VRAM Optimization: Moves model to GPU only during encoding, then offloads back to CPU to free VRAM for diffusion model. Args: tokens: Dict with 'input_ids' and 'attention_mask' tensors Returns: Tuple of (hidden_states, pooled_output, extra_dict) where extra_dict contains 'attention_mask' for the diffusion model """ input_ids = tokens.get("input_ids") if input_ids is None: raise ValueError("tokens dict must contain 'input_ids'") # Move model to GPU for encoding logger.info(f"Moving text encoder to {self.device} for encoding...") self.model = self.model.to(self.device) input_ids = input_ids.to(self.device) # Get attention mask if present - CRITICAL for proper masking of padding tokens attention_mask = tokens.get("attention_mask") if attention_mask is not None: attention_mask = attention_mask.to(self.device) try: with torch.no_grad(): outputs = self.model(input_ids, attention_mask=attention_mask) # Return concatenated hidden states, pooled output, and extra with attention_mask hidden_states = outputs["hidden_states"].clone() # Clone to keep on GPU pooled_out = outputs["pooled_output"] pooled = pooled_out.clone() if pooled_out is not None else None # Clone if exists finally: # Offload model back to CPU to free VRAM for diffusion model logger.info(f"Offloading text encoder to {self.offload_device}...") self.model = self.model.to(self.offload_device) if torch.cuda.is_available(): torch.cuda.empty_cache() # Return attention mask in extra dict for the diffusion model to use extra = {} if attention_mask is not None: extra["attention_mask"] = attention_mask return hidden_states, pooled, extra def tokenize(self, text, return_word_ids=False): """Tokenize text (CLIP interface compatibility for Adetailer). Args: text: Text to tokenize return_word_ids: Whether to return word IDs Returns: Dict with 'input_ids' and 'attention_mask' """ return self.tokenizer.tokenize_with_weights(text, return_word_ids) def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False): """Encode from tokens (CLIP interface compatibility for Adetailer). Args: tokens: Dict with 'input_ids' and 'attention_mask' return_pooled: Whether to return pooled output return_dict: Whether to return as dict Returns: Embeddings, or (embeddings, pooled) if return_pooled, or dict if return_dict """ cond, pooled, extra = self.encode_token_weights(tokens) if return_dict: out = {"cond": cond, "pooled_output": pooled} out.update(extra) return out return (cond, pooled) if return_pooled else cond def load_model(self): """Load model to GPU (CLIP interface compatibility). Returns: Self for compatibility """ # Move model to device if not already there if self.device is not None: self.model = self.model.to(self.device) return self def load_sd(self, state_dict: dict) -> tuple: """Load state dictionary into model. Args: state_dict: Model weights Returns: Tuple of (missing_keys, unexpected_keys) """ # Filter and map state dict keys for Qwen3 model model_sd = {} for k, v in state_dict.items(): # Map state dict keys to model structure if k.startswith("model."): model_sd[k[6:]] = v # Remove "model." prefix else: model_sd[k] = v return self.model.load_state_dict(model_sd, strict=False) def klein_clip(dtype=None) -> dict: """Create Klein CLIP configuration. Returns: Dict with 'clip' and 'tokenizer' classes """ class Target: clip = KleinCLIP tokenizer = KleinTokenizer params = {"dtype": dtype} return Target # Convenience function to detect Klein model from state dict def detect_klein_model(state_dict: dict) -> bool: """Detect if state dict is from a Klein text encoder. Args: state_dict: Model state dictionary Returns: True if this appears to be a Klein model """ klein_indicators = [ "model.layers.0.self_attn.q_norm.weight", "model.layers.0.self_attn.k_norm.weight", "embed_tokens.weight", ] keys = set(state_dict.keys()) for indicator in klein_indicators: for key in keys: if indicator in key: return True return False