From 4ea0cac77419aa074e19da81d3ad0ea774d6b175 Mon Sep 17 00:00:00 2001
From: lodestone-rock <lodestone.rock@gmail.com>
Date: Sun, 22 Jun 2025 16:44:05 +0700
Subject: [PATCH] init flux lite distillation

---
 flux_mod/loader.py | 121 ++++++++++++++++++++-
 flux_mod/model.py  | 254 +++++++++++++++++++++++++++++++++++++++++++++
 flux_mod/nodes.py  |  63 ++++++++++-
 3 files changed, 436 insertions(+), 2 deletions(-)

diff --git a/flux_mod/loader.py b/flux_mod/loader.py
index f12813a..8ac7b01 100644
--- a/flux_mod/loader.py
+++ b/flux_mod/loader.py
@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from safetensors.torch import safe_open
-from .model import FluxParams, FluxMod
+from .model import FluxParams, FluxMod, FluxLiteMod
 from .layers import Approximator
 import comfy
 import comfy.model_patcher
@@ -250,3 +250,122 @@ def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.b
         offload_device = offload_device,
     )
     return model_patcher
+
+
+def load_flux_lite_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.bfloat16, lite_patch_path=None, is_gguf=False):
+
+    if is_gguf:
+        ensure_gguf()
+    # just load safetensors here
+    state_dict = load_selected_keys(model_path, {"mod", "time_in", "guidance_in", "vector_in"}, is_gguf=is_gguf)
+
+    if timestep_guidance_path is None:
+        # Chroma mode - we expect the timestep guidance to be bundled under
+        # the key distilled_guidance_layer.
+        if lite_patch_path is not None:
+            raise ValueError("Internal error: lite patch specified in Chroma loader mode")
+        timestep_state_dict = {
+            k[25:]: v
+            for k, v in state_dict.items()
+            if k.startswith("distilled_guidance_layer.")
+        }
+        if not timestep_state_dict:
+            raise RuntimeError("Could not find distilled guidance layer in Chroma model")
+        n_layers = 0
+        for key in timestep_state_dict:
+            keysplit = key.split(".", 2)
+            if (
+                len(keysplit) == 3 and
+                keysplit[0] == "norms" and
+                keysplit[1].isnumeric() and
+                keysplit[2] == "scale"
+            ):
+                n_layers = max(n_layers, int(keysplit[1]) + 1)
+            del state_dict[f"distilled_guidance_layer.{key}"]
+        if n_layers == 0:
+            raise RuntimeError("Could not determine number of distilled guidance layers in Chroma model")
+    else:
+        timestep_state_dict = comfy.utils.load_torch_file(timestep_guidance_path)
+
+        n_layers = 5
+        # if "v3" in timestep_guidance_path:
+        #     n_layers = 6
+        # elif "v2" in timestep_guidance_path:
+        #     n_layers = 5
+        # else:
+        #     n_layers = 4
+
+    param_count = sum(x.numel() for x in state_dict.values())
+    unet_dtype = torch.bfloat16
+
+    load_device = model_management.get_torch_device()
+    offload_device = model_management.unet_offload_device()
+
+
+    params=FluxParams(
+        in_channels=64,
+        vec_in_dim=768,
+        context_in_dim=4096,
+        hidden_size=3072,
+        mlp_ratio=4.0,
+        num_heads=24,
+        depth=8,
+        depth_single_blocks=38,
+        axes_dim=[16, 56, 56],
+        theta=10_000,
+        qkv_bias=True,
+        guidance_embed=False,
+    )
+
+    model_conf = ExternalFlux()
+    model_class = ChromaFluxModel if timestep_guidance_path is None else ExternalFluxModel
+    model = model_class(
+        model_conf,
+        model_type=comfy.model_base.ModelType.FLUX,
+        device=model_management.get_torch_device()
+    )
+    unet_config = model_conf.unet_config
+    if is_gguf:
+        operations = model_conf.custom_operations = gguf.ops.GGMLOps()
+        modelpatcher_class = gguf.nodes.GGUFModelPatcher
+    else:
+        modelpatcher_class = comfy.model_patcher.ModelPatcher
+    if model_conf.custom_operations is None:
+        fp8 = model_conf.optimizations.get("fp8", model_conf.scaled_fp8 is not None)
+        operations = comfy.ops.pick_operations(
+            unet_config.get("dtype"),
+            model.manual_cast_dtype,
+            fp8_optimizations=fp8,
+            scaled_fp8=model_conf.scaled_fp8,
+        )
+    else:
+        operations = model_conf.custom_operations
+    
+    model.diffusion_model = FluxLiteMod(params=params, dtype=unet_dtype, operations=operations)
+
+    model.diffusion_model.load_state_dict(state_dict)
+    
+    # if lite_patch_path is not None:
+    #     model.diffusion_model.lite = True
+        # for _ in range(5,16):
+        #     del model.diffusion_model.double_blocks[5]
+
+        # model.diffusion_model.double_blocks[4].load_state_dict(comfy.utils.load_torch_file(lite_patch_path))
+
+    model.diffusion_model.distilled_guidance_layer = Approximator(32, 3072, 5120, n_layers, operations=operations)
+    model.diffusion_model.distilled_guidance_layer.load_state_dict(timestep_state_dict)
+    model.diffusion_model.dtype = unet_dtype
+    model.diffusion_model.eval()
+    if not is_gguf:
+        model.diffusion_model.to(unet_dtype)
+        # we cast to fp8 for mixed matmul ops but omit the picky and sensitive layers
+        cast_layers(model.diffusion_model, nn.Linear, dtype=linear_dtypes, exclude_keywords={"img_in", "final_layer", "scale"})
+        if model_management.force_channels_last():
+            model.diffusion_model.to(memory_format=torch.channels_last)
+
+    model_patcher = modelpatcher_class(
+        model,
+        load_device = load_device,
+        offload_device = offload_device,
+    )
+    return model_patcher
\ No newline at end of file
diff --git a/flux_mod/model.py b/flux_mod/model.py
index 31dc1d9..16bf11a 100644
--- a/flux_mod/model.py
+++ b/flux_mod/model.py
@@ -326,3 +326,257 @@ class FluxMod(nn.Module):
         txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
         out = self.forward_orig(img, img_ids, context, txt_ids, timestep, guidance, control, transformer_options)
         return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
+
+
+class FluxLiteMod(nn.Module):
+    """
+    Transformer model for flow matching on sequences.
+    """
+
+    def __init__(self, params: FluxParams, image_model=None, final_layer=True, dtype=None, device=None, operations=nn, **kwargs):
+        super().__init__()
+        self.dtype = dtype
+        self.params = params
+        self.in_channels = params.in_channels
+        self.out_channels = self.in_channels
+        if params.hidden_size % params.num_heads != 0:
+            raise ValueError(
+                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+            )
+        pe_dim = params.hidden_size // params.num_heads
+        if sum(params.axes_dim) != pe_dim:
+            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+        self.hidden_size = params.hidden_size
+        self.num_heads = params.num_heads
+        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+        self.img_in = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device)
+        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
+        # set as nn identity for now, will overwrite it later.
+        self.distilled_guidance_layer = nn.Identity()
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamBlock(
+                    self.hidden_size,
+                    self.num_heads,
+                    mlp_ratio=params.mlp_ratio,
+                    qkv_bias=params.qkv_bias,
+                    dtype=dtype, device=device, operations=operations
+                )
+                for _ in range(params.depth)
+            ]
+        )
+
+        self.single_blocks = nn.ModuleList(
+            [
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, dtype=dtype, device=device, operations=operations)
+                for _ in range(params.depth_single_blocks)
+            ]
+        )
+
+        if final_layer:
+            self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
+
+        self.skip_mmdit = []
+        self.skip_dit = []
+
+    @staticmethod
+    def distribute_modulations(tensor: torch.Tensor, single_block_count: int = 38, double_blocks_count: int = 8):
+        """
+        Distributes slices of the tensor into the block_dict as ModulationOut objects.
+
+        Args:
+            tensor (torch.Tensor): Input tensor with shape [batch_size, vectors, dim].
+        """
+        batch_size, vectors, dim = tensor.shape
+
+        block_dict = {}
+
+        # HARD CODED VALUES! lookup table for the generated vectors
+        # Add 38 single mod blocks
+        for i in range(single_block_count):
+            key = f"single_blocks.{i}.modulation.lin"
+            block_dict[key] = None
+
+        # Add 19 image double blocks
+        for i in range(double_blocks_count):
+            key = f"double_blocks.{i}.img_mod.lin"
+            block_dict[key] = None
+
+        # Add 19 text double blocks
+        for i in range(double_blocks_count):
+            key = f"double_blocks.{i}.txt_mod.lin"
+            block_dict[key] = None
+
+        # Add the final layer
+        block_dict["final_layer.adaLN_modulation.1"] = None
+
+
+        idx = 0  # Index to keep track of the vector slices
+
+        for key in block_dict.keys():
+            if "single_blocks" in key:
+                # Single block: 1 ModulationOut
+                block_dict[key] = ModulationOut(
+                    shift=tensor[:, idx:idx+1, :],
+                    scale=tensor[:, idx+1:idx+2, :],
+                    gate=tensor[:, idx+2:idx+3, :]
+                )
+                idx += 3  # Advance by 3 vectors
+
+            elif "img_mod" in key:
+                # Double block: List of 2 ModulationOut
+                double_block = []
+                for _ in range(2):  # Create 2 ModulationOut objects
+                    double_block.append(
+                        ModulationOut(
+                            shift=tensor[:, idx:idx+1, :],
+                            scale=tensor[:, idx+1:idx+2, :],
+                            gate=tensor[:, idx+2:idx+3, :]
+                        )
+                    )
+                    idx += 3  # Advance by 3 vectors per ModulationOut
+                block_dict[key] = double_block
+
+            elif "txt_mod" in key:
+                # Double block: List of 2 ModulationOut
+                double_block = []
+                for _ in range(2):  # Create 2 ModulationOut objects
+                    double_block.append(
+                        ModulationOut(
+                            shift=tensor[:, idx:idx+1, :],
+                            scale=tensor[:, idx+1:idx+2, :],
+                            gate=tensor[:, idx+2:idx+3, :]
+                        )
+                    )
+                    idx += 3  # Advance by 3 vectors per ModulationOut
+                block_dict[key] = double_block
+
+            elif "final_layer" in key:
+                # Final layer: 1 ModulationOut
+                block_dict[key] = [
+                    tensor[:, idx:idx+1, :],
+                    tensor[:, idx+1:idx+2, :],
+                ]
+                idx += 2  # Advance by 2 vectors
+
+        return block_dict
+
+    def forward_orig(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        guidance: Tensor = None,
+        control=None,
+        transformer_options={},
+    ) -> Tensor:
+        patches_replace = transformer_options.get("patches_replace", {})
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+
+        # distilled vector guidance
+        device = img.device
+        dtype = img.dtype
+        mod_index_length = 212
+        distill_timestep = timestep_embedding(timesteps.detach().clone(), 8).to(device=device, dtype=dtype)
+        # guidance = guidance *
+        distil_guidance = timestep_embedding(guidance.detach().clone(), 8).to(device=device, dtype=dtype)
+
+        # get all modulation index
+        modulation_index = timestep_embedding(torch.arange(mod_index_length), 16).to(device=device, dtype=dtype)
+        # we need to broadcast the modulation index here so each batch has all of the index
+        modulation_index = modulation_index.unsqueeze(0).repeat(img.shape[0], 1, 1)
+        # and we need to broadcast timestep and guidance along too
+        timestep_guidance = torch.cat([distill_timestep, distil_guidance], dim=1).unsqueeze(1).repeat(1, mod_index_length, 1)
+        # then and only then we could concatenate it together
+        input_vec = torch.cat([timestep_guidance, modulation_index], dim=-1)
+
+        mod_vectors = self.distilled_guidance_layer(input_vec)
+
+        mod_vectors_dict = self.distribute_modulations(mod_vectors, 38, 8)
+
+
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        blocks_replace = patches_replace.get("dit", {})
+        for i, block in enumerate(self.double_blocks):
+            if i not in self.skip_mmdit:
+                guidance_index = i
+
+                img_mod = mod_vectors_dict[f"double_blocks.{guidance_index}.img_mod.lin"]
+                txt_mod = mod_vectors_dict[f"double_blocks.{guidance_index}.txt_mod.lin"]
+                double_mod = [img_mod, txt_mod]
+
+                if ("double_block", i) in blocks_replace:
+                    def block_wrap(args):
+                        out = {}
+                        out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], pe=args["pe"], distill_vec=args["distill_vec"])
+                        return out
+
+                    out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "pe": pe, "distill_vec": double_mod}, {"original_block": block_wrap})
+                    txt = out["txt"]
+                    img = out["img"]
+                else:
+                    img, txt = block(img=img, txt=txt, pe=pe, distill_vec=double_mod)
+
+                if control is not None: # Controlnet
+                    control_i = control.get("input")
+                    if i < len(control_i):
+                        add = control_i[i]
+                        if add is not None:
+                            img += add
+
+        img = torch.cat((txt, img), 1)
+
+        for i, block in enumerate(self.single_blocks):
+            if i not in self.skip_dit:
+                single_mod = mod_vectors_dict[f"single_blocks.{i}.modulation.lin"]
+                if ("single_block", i) in blocks_replace:
+
+                    def block_wrap(args):
+                        out = {}
+                        out["img"] = block(args["img"], pe=args["pe"], distill_vec=args["distill_vec"])
+                        return out
+
+                    out = blocks_replace[("single_block", i)]({"img": img, "pe": pe, "distill_vec": single_mod}, {"original_block": block_wrap})
+                    img = out["img"]
+                else:
+                    img = block(img, pe=pe, distill_vec=single_mod)
+
+                if control is not None: # Controlnet
+                    control_o = control.get("output")
+                    if i < len(control_o):
+                        add = control_o[i]
+                        if add is not None:
+                            img[:, txt.shape[1] :, ...] += add
+
+        img = img[:, txt.shape[1] :, ...]
+        final_mod = mod_vectors_dict["final_layer.adaLN_modulation.1"]
+        img = self.final_layer(img, distill_vec=final_mod)  # (N, T, patch_size ** 2 * out_channels)
+        return img
+
+    def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
+        bs, c, h, w = x.shape
+        patch_size = 2
+        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
+
+        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+
+        h_len = ((h + (patch_size // 2)) // patch_size)
+        w_len = ((w + (patch_size // 2)) // patch_size)
+        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
+        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+        out = self.forward_orig(img, img_ids, context, txt_ids, timestep, guidance, control, transformer_options)
+        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
diff --git a/flux_mod/nodes.py b/flux_mod/nodes.py
index 4996f9c..01cec36 100644
--- a/flux_mod/nodes.py
+++ b/flux_mod/nodes.py
@@ -4,7 +4,7 @@ import json
 import torch
 import folder_paths
 
-from .loader import load_flux_mod
+from .loader import load_flux_mod, load_flux_lite_mod
 from .sampler import common_ksampler
 import comfy.samplers
 import comfy.cli_args
@@ -79,6 +79,66 @@ class FluxModDiffusionLoader:
         return (flux_mod,)
 
 
+class FluxLiteModDiffusionLoader:
+    NodeId = "FluxLiteModDiffusionLoader"
+    NodeName = "FluxLiteMod Unified Model Loader"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        patches = ["None"] + folder_paths.get_filename_list("diffusion_models")
+        checkpoint_paths = folder_paths.get_filename_list("diffusion_models")
+        if "unet_gguf" in folder_paths.folder_names_and_paths:
+            checkpoint_paths = checkpoint_paths + folder_paths.get_filename_list(
+                "unet_gguf"
+            )
+        return {
+            "required": {
+                "unet_name": (checkpoint_paths, {"tooltip": "This loads both .safetensors and .GGUF models if ComfyUI-GGUF is installed."}),
+                "guidance_name": (patches,),
+                "quant_mode": (
+                    ["bf16", "float8_e4m3fn (8 bit)", "float8_e5m2 (also 8 bit)"],
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_unet"
+    CATEGORY = "advanced/FluxMod"
+
+    def load_unet(
+        self, *, unet_name, quant_mode, guidance_name=None, lite_patch_unet_name=None
+    ):
+        dtypes = {
+            "bf16": torch.bfloat16,
+            "float8_e4m3fn (8 bit)": torch.float8_e4m3fn,
+            "float8_e5m2 (also 8 bit)": torch.float8_e5m2,
+        }
+
+        is_gguf = unet_name.lower().endswith(".gguf")
+        unet_path = folder_paths.get_full_path(
+            "unet_gguf" if is_gguf else "diffusion_models", unet_name
+        )
+        if guidance_name is not None:
+            guidance_path = folder_paths.get_full_path(
+                "diffusion_models", guidance_name
+            )
+        else:
+            guidance_path = None
+        if lite_patch_unet_name is not None:
+            lite_patch_unet_name = folder_paths.get_full_path(
+                "diffusion_models", lite_patch_unet_name
+            )
+        flux_mod = load_flux_lite_mod(
+            model_path=unet_path,
+            timestep_guidance_path=guidance_path,
+            linear_dtypes=dtypes[quant_mode],
+            lite_patch_path=lite_patch_unet_name,
+            is_gguf=is_gguf,
+        )
+        return (flux_mod,)
+    
+
 class FluxModDiffusionLoaderMini(FluxModDiffusionLoader):
     NodeId = "FluxModDiffusionLoaderMini"
     NodeName = "FluxMod Mini Unified Model Loader"
@@ -420,6 +480,7 @@ class FluxModSamplerWrapperNode:
 node_list = [
     # loaders
     FluxModDiffusionLoader,
+    FluxLiteModDiffusionLoader,
     FluxModDiffusionLoaderMini,
     ChromaDiffusionLoader,
     # Sampler