fancyfeast/

flux2klein-optimized-text-embedder-4B

from __future__ import annotations

import json
import math
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F
from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
from transformers import PreTrainedModel


class FluxQwen3TorchEmbedder(PreTrainedModel):
	"""Stripped down and optimized Qwen3 specifically for Flux 2 Klein models.
	In my testing this is about 1.3x faster than using the original HF implementation, and saves ~3GB of peak memory on the 8GB model.

	The output_hidden_state_indices is 9, 18, 27 for both Klein 4B and Klein 9B.
	"""
	config_class = Qwen3Config
	base_model_prefix = "flux_qwen3"

	def __init__(
		self,
		config: Qwen3Config,
		*,
		output_hidden_state_indices: tuple[int, ...] = (9, 18, 27),
		max_sequence_length: int = 512,
	):
		super().__init__(config)

		self.hidden_size = config.hidden_size
		self.num_attention_heads = config.num_attention_heads
		self.head_dim = int(getattr(config, "head_dim", self.hidden_size // self.num_attention_heads))
		self.rope_theta = float(getattr(config, "rope_theta", 1000000.0))

		self.output_hidden_state_indices = tuple(int(i) for i in output_hidden_state_indices)
		if not self.output_hidden_state_indices:
			raise ValueError("output_hidden_state_indices must not be empty")
		if min(self.output_hidden_state_indices) < 1:
			raise ValueError("output hidden state indices must be >= 1 for decoder layer outputs")
		
		max_layer_needed = max(self.output_hidden_state_indices)
		if max_layer_needed > int(config.num_hidden_layers):
			raise ValueError(f"requested hidden state after layer {max_layer_needed}, but config.num_hidden_layers={config.num_hidden_layers}")

		self.capture_slot_by_layer = {
			layer_idx: slot for slot, layer_idx in enumerate(self.output_hidden_state_indices)
		}
		self.max_sequence_length = int(max_sequence_length)

		self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=getattr(config, "pad_token_id", None))
		self.layers = nn.ModuleList(
			FluxQwen3TorchLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)
		)

		# Built lazily/refreshed in forward so dtype/device tracks the model.
		self.register_buffer("cos_cached", torch.empty(0), persistent=False)
		self.register_buffer("sin_cached", torch.empty(0), persistent=False)
		self.register_buffer("causal_mask", torch.empty(0, dtype=torch.bool), persistent=False)

		self.post_init()
	
	def _maybe_refresh_caches(self, *, device: torch.device, dtype: torch.dtype):
		need_refresh = self.cos_cached.numel() == 0 or self.cos_cached.shape[0] < self.max_sequence_length or self.cos_cached.device != device or self.cos_cached.dtype != dtype
		if not need_refresh:
			return
		
		cos, sin = _rotary_cache(
			self.max_sequence_length,
			self.head_dim,
			self.rope_theta,
			device=device,
			dtype=dtype,
		)

		pos = torch.arange(self.max_sequence_length, device=device)
		causal = pos[None, :] <= pos[:, None]

		self.cos_cached = cos
		self.sin_cached = sin
		self.causal_mask = causal[None, None, :, :]

	@classmethod
	def _from_original_hf_checkpoint(cls, checkpoint_path: str, subfolder: str | None) -> "FluxQwen3TorchEmbedder":
		from huggingface_hub import hf_hub_download
		import safetensors.torch
		from transformers import AutoConfig

		cfg = AutoConfig.from_pretrained(checkpoint_path, subfolder=subfolder)
		assert isinstance(cfg, Qwen3Config), f"expected Qwen3Config, got {type(cfg)}"
		cfg.num_hidden_layers = 27
		if cfg.layer_types is not None:
			cfg.layer_types = cfg.layer_types[:27]
		cfg.max_window_layers = 27
		model = cls(cfg)

		# Load the original checkpoint
		index_path = hf_hub_download(checkpoint_path, filename="model.safetensors.index.json", subfolder=subfolder)
		index = json.loads(Path(index_path).read_text())
		shard_names = set(index['weight_map'].values())
		original_checkpoint = {}

		for shard_name in shard_names:
			path = hf_hub_download(checkpoint_path, filename=shard_name, subfolder=subfolder)
			shard = safetensors.torch.load_file(path)
			original_checkpoint.update(shard)

		# Copy weights from the original checkpoint into our model
		with torch.no_grad():
			model.embed_tokens.weight.copy_(original_checkpoint["model.embed_tokens.weight"])

			for layer_idx in range(len(model.layers)):
				layer = model.layers[layer_idx]
				layer_base = f"model.layers.{layer_idx}."

				layer.input_layernorm_weight.copy_(original_checkpoint[layer_base + "input_layernorm.weight"])
				layer.post_attention_layernorm_weight.copy_(original_checkpoint[layer_base + "post_attention_layernorm.weight"])
				q = original_checkpoint[layer_base + "self_attn.q_proj.weight"]
				k = original_checkpoint[layer_base + "self_attn.k_proj.weight"]
				v = original_checkpoint[layer_base + "self_attn.v_proj.weight"]
				layer.qkv_proj_weight.copy_(torch.cat((q, k, v), dim=0))
				layer.o_proj_weight.copy_(original_checkpoint[layer_base + "self_attn.o_proj.weight"])
				layer.q_norm_weight.copy_(original_checkpoint[layer_base + "self_attn.q_norm.weight"])
				layer.k_norm_weight.copy_(original_checkpoint[layer_base + "self_attn.k_norm.weight"])
				gate = original_checkpoint[layer_base + "mlp.gate_proj.weight"]
				up = original_checkpoint[layer_base + "mlp.up_proj.weight"]
				layer.gate_up_proj_weight.copy_(torch.cat((gate, up), dim=0))
				layer.down_proj_weight.copy_(original_checkpoint[layer_base + "mlp.down_proj.weight"])
		
		return model

	def forward(
		self,
		input_ids: torch.Tensor,
		attention_mask: torch.Tensor,
	) -> torch.Tensor:
		if input_ids.ndim != 2:
			raise ValueError(f"expected input_ids [batch, seq], got {tuple(input_ids.shape)}")
		
		batch, seq_len = input_ids.shape

		if seq_len != self.max_sequence_length:
			raise ValueError(f"sequence length {seq_len} does not match cached max {self.max_sequence_length}")

		dtype = self.embed_tokens.weight.dtype
		device = input_ids.device
		self._maybe_refresh_caches(device=device, dtype=dtype)

		key_mask = attention_mask.reshape(batch, 1, 1, seq_len).to(dtype=torch.bool)
		sdpa_mask = self.causal_mask[:, :, :seq_len, :seq_len] & key_mask

		cos = self.cos_cached[:seq_len]
		sin = self.sin_cached[:seq_len]

		hidden_states = self.embed_tokens(input_ids)

		prompt_embeds = torch.empty(
			batch,
			seq_len,
			len(self.output_hidden_state_indices) * self.hidden_size,
			device=input_ids.device,
			dtype=dtype,
		)

		for layer_number, layer in enumerate(self.layers, start=1):
			hidden_states = layer(hidden_states, cos, sin, sdpa_mask)
			slot = self.capture_slot_by_layer.get(layer_number)
			if slot is None:
				continue

			start = slot * self.hidden_size
			prompt_embeds[:, :, start : start + self.hidden_size].copy_(hidden_states)

		return prompt_embeds


class FluxQwen3TorchLayer(nn.Module):
	def __init__(self, config: Qwen3Config, layer_idx: int):
		super().__init__()

		self.layer_idx = layer_idx
		self.hidden_size = config.hidden_size
		self.intermediate_size = config.intermediate_size
		self.num_attention_heads = config.num_attention_heads
		assert config.num_key_value_heads is not None, "num_key_value_heads must be specified in config for FluxQwen3TorchLayer"
		self.num_key_value_heads = config.num_key_value_heads
		self.head_dim = int(getattr(config, "head_dim", self.hidden_size // self.num_attention_heads))
		self.rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-6))
		self.scale = 1.0 / math.sqrt(self.head_dim)

		self.q_width = self.num_attention_heads * self.head_dim
		self.kv_width = self.num_key_value_heads * self.head_dim
		self.k_offset = self.q_width
		self.v_offset = self.q_width + self.kv_width

		self.input_layernorm_weight = nn.Parameter(torch.empty(self.hidden_size))
		self.post_attention_layernorm_weight = nn.Parameter(torch.empty(self.hidden_size))

		self.qkv_proj_weight = nn.Parameter(torch.empty(self.q_width + 2 * self.kv_width, self.hidden_size))
		self.o_proj_weight = nn.Parameter(torch.empty(self.hidden_size, self.q_width))

		self.q_norm_weight = nn.Parameter(torch.empty(self.head_dim))
		self.k_norm_weight = nn.Parameter(torch.empty(self.head_dim))

		self.gate_up_proj_weight = nn.Parameter(torch.empty(self.intermediate_size * 2, self.hidden_size))
		self.down_proj_weight = nn.Parameter(torch.empty(self.hidden_size, self.intermediate_size))

		assert self.q_width == self.o_proj_weight.shape[1]
		assert self.o_proj_weight.shape == (self.hidden_size, self.q_width)
		assert self.qkv_proj_weight.shape == (self.q_width + 2 * self.kv_width, self.hidden_size)

	def _rms_norm(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
		dtype = x.dtype
		x_float = x.float()
		variance = x_float.pow(2).mean(dim=-1, keepdim=True)
		return (x_float * torch.rsqrt(variance + self.rms_norm_eps)).to(dtype) * weight

	def _head_rms_norm(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
		dtype = x.dtype
		x_float = x.float()
		variance = x_float.pow(2).mean(dim=-1, keepdim=True)
		return (x_float * torch.rsqrt(variance + self.rms_norm_eps)).to(dtype) * weight

	@staticmethod
	def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
		half = x.shape[-1] // 2
		x1 = x[..., :half]
		x2 = x[..., half:]
		rotated = torch.cat((-x2, x1), dim=-1)
		return x * cos[:, None, :] + rotated * sin[:, None, :]

	def forward(
		self,
		hidden_states: torch.Tensor,
		cos: torch.Tensor,
		sin: torch.Tensor,
		attention_mask: torch.Tensor,
	) -> torch.Tensor:
		residual = hidden_states
		x = self._rms_norm(hidden_states, self.input_layernorm_weight)

		batch, seq_len, _ = x.shape
		qkv = F.linear(x, self.qkv_proj_weight)
		q_raw = qkv[:, :, : self.q_width].view(batch, seq_len, self.num_attention_heads, self.head_dim)
		k_raw = qkv[:, :, self.k_offset : self.v_offset].view(
			batch, seq_len, self.num_key_value_heads, self.head_dim
		)
		v = qkv[:, :, self.v_offset :].view(batch, seq_len, self.num_key_value_heads, self.head_dim)

		q = self._apply_rope(self._head_rms_norm(q_raw, self.q_norm_weight), cos, sin).transpose(1, 2)
		k = self._apply_rope(self._head_rms_norm(k_raw, self.k_norm_weight), cos, sin).transpose(1, 2)
		v = v.transpose(1, 2)

		attn = F.scaled_dot_product_attention(
			q,
			k,
			v,
			attn_mask=attention_mask,
			dropout_p=0.0,
			scale=self.scale,
			is_causal=False,
			enable_gqa=True,
		)
		attn = attn.transpose(1, 2).contiguous().view(batch, seq_len, self.q_width)
		hidden_states = residual + F.linear(attn, self.o_proj_weight)

		residual = hidden_states
		x = self._rms_norm(hidden_states, self.post_attention_layernorm_weight)

		gate_up = F.linear(x, self.gate_up_proj_weight)
		gate, up = gate_up.split(self.intermediate_size, dim=-1)
		x = F.silu(gate) * up

		hidden_states = residual + F.linear(x, self.down_proj_weight)
		return hidden_states


def _rotary_cache(
	seq_len: int,
	head_dim: int,
	rope_theta: float,
	*,
	device: torch.device | str,
	dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
	inv_freq = 1.0 / (
		rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
	)
	pos = torch.arange(seq_len, dtype=torch.float32, device=device)
	freqs = torch.outer(pos, inv_freq)
	emb = torch.cat((freqs, freqs), dim=-1)
	return emb.cos().to(dtype=dtype).contiguous(), emb.sin().to(dtype=dtype).contiguous()