From d4ffbcad708e23142a479f08b5ce95d35c46c41e Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Wed, 14 Oct 2026 13:03:00 -0600 Subject: [PATCH] flux3: drop dead gate_type path, reject flow-map packages, note refreshed base checkpoint --- .../diffusion_models/flux3/flux3.py | 5 +++++ .../diffusion_models/flux3/src/transformer.py | 17 ----------------- 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/extensions_built_in/diffusion_models/flux3/flux3.py b/extensions_built_in/diffusion_models/flux3/flux3.py index effb320..1a40dc5 100644 --- a/extensions_built_in/diffusion_models/flux3/flux3.py +++ b/extensions_built_in/diffusion_models/flux3/flux3.py @@ -193,6 +193,11 @@ class Flux3Model(BaseModel): f"({cfg['quantization']}) are not supported; use the bf16 package " "and let the toolkit quantize." ) + if cfg is not None and cfg["dit"].get("flow_map", False): + raise ValueError( + "FLUX.3 step-distilled (flow-map) packages are not supported; " + "use the base package." + ) state_dict = load_file(dit_path) state_dict = { diff --git a/extensions_built_in/diffusion_models/flux3/src/transformer.py b/extensions_built_in/diffusion_models/flux3/src/transformer.py index 3a6c640..54db610 100644 --- a/extensions_built_in/diffusion_models/flux3/src/transformer.py +++ b/extensions_built_in/diffusion_models/flux3/src/transformer.py @@ -66,7 +66,6 @@ class Flux3TransformerParams: theta: int = 10000 mlp_ratio: float = 3.0 qkv_bias: bool = False - gate_type: str | None = None attn_mode: str = "native" # modes with a final_layer (velocity outputs); None = every content mode out_modes: list[str] | None = field(default_factory=lambda: ["image"]) @@ -248,7 +247,6 @@ class ModeBlock(nn.Module): num_heads: int, mlp_ratio: float = 3.0, qkv_bias: bool = False, - gate_type: str | None = None, use_gated_mlp: bool = True, ): super().__init__() @@ -257,7 +255,6 @@ class ModeBlock(nn.Module): self.hidden_size = hidden_size self.mlp_hidden_dim = int(hidden_size * mlp_ratio) self.mlp_mult_factor = 2 if use_gated_mlp else 1 - self.gate_type = gate_type self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) @@ -265,8 +262,6 @@ class ModeBlock(nn.Module): self.mlp_in = nn.Linear( hidden_size, self.mlp_hidden_dim * self.mlp_mult_factor, bias=False ) - if gate_type is not None: - self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False) self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) @@ -307,8 +302,6 @@ class ModeBlock(nn.Module): output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out( _bf16_if_required(self.mlp_out, self.mlp_act(mlp)) ) - if self.gate_type is not None: - output = output * F.silu(self.gate_proj(x_mod)) return x + mod_gate * output @@ -320,7 +313,6 @@ class SingleStreamBlock(nn.Module): hidden_size: int, num_heads: int, mlp_ratio: float = 3.0, - gate_type: str | None = None, use_gated_mlp: bool = True, ): super().__init__() @@ -329,7 +321,6 @@ class SingleStreamBlock(nn.Module): self.hidden_size = hidden_size self.mlp_hidden_dim = int(hidden_size * mlp_ratio) self.mlp_mult_factor = 2 if use_gated_mlp else 1 - self.gate_type = gate_type self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False) @@ -339,8 +330,6 @@ class SingleStreamBlock(nn.Module): ) self.attn_out = nn.Linear(hidden_size, hidden_size, bias=False) self.mlp_out = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) - if gate_type is not None: - self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.norm = QKNorm(head_dim) self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) @@ -389,8 +378,6 @@ class SingleStreamBlock(nn.Module): output = self.attn_out(_bf16_if_required(self.attn_out, attn)) + self.mlp_out( _bf16_if_required(self.mlp_out, self.mlp_act(mlp)) ) - if self.gate_type is not None: - output = output * F.silu(self.gate_proj(combined)) ctx_out, x_out = torch.split(output, [ctx.shape[1], x.shape[1]], dim=1) return ctx + mod_gate_ctx * ctx_out, x + mod_gate_x * x_out @@ -428,7 +415,6 @@ class Flux3Transformer(nn.Module, OstrisModelMixin): self.depth_single_blocks = params.depth_single_blocks self.mlp_ratio = params.mlp_ratio self.qkv_bias = params.qkv_bias - self.gate_type = params.gate_type self.attn_mode = params.attn_mode self.gradient_checkpointing = False @@ -494,7 +480,6 @@ class Flux3Transformer(nn.Module, OstrisModelMixin): params.num_heads, mlp_ratio=params.mlp_ratio, qkv_bias=params.qkv_bias, - gate_type=params.gate_type, ) for _ in range(params.depth) ] @@ -509,7 +494,6 @@ class Flux3Transformer(nn.Module, OstrisModelMixin): params.num_heads, mlp_ratio=params.mlp_ratio, qkv_bias=params.qkv_bias, - gate_type=params.gate_type, ) for _ in range(params.depth) ] @@ -521,7 +505,6 @@ class Flux3Transformer(nn.Module, OstrisModelMixin): params.hidden_size, params.num_heads, mlp_ratio=params.mlp_ratio, - gate_type=params.gate_type, ) for _ in range(params.depth_single_blocks) ]