Inline kernel setup: Diffusers downloads only each component's own module

NOTICE

@@ -30,4 +30,4 @@ MODIFIED FILES (changed by joseplcam):
   The `text_encoder` and `transformer` entries point to the custom classes below.
 
 NEW FILES: text_encoder/modeling_nunchaku_qwen3vl.py, transformer/modeling_nunchaku_qwenimage21.py,
-text_encoder/nunchaku_kernels.py, transformer/nunchaku_kernels.py.
+tools/ (the scripts that produced this repository).

README.md

@@ -48,11 +48,13 @@ edited = pipe("Make it night time with moonlight", image=image, output_resolutio
 
 - `text_encoder/modeling_nunchaku_qwen3vl.py`: a `Qwen3VLForConditionalGeneration` subclass that swaps the
   quantized linears for Diffusers' own `SVDQW4A4Linear` before loading the weights.
-- `nunchaku_kernels.py` (in `text_encoder/` and `transformer/`): Diffusers loads the NVFP4 kernels from
-  `rootonchair/nunchaku-lite-kernels`, which is no longer downloadable. This file points that name at
-  [joseplcam/nunchaku-lite-kernels](https://huggingface.co/joseplcam/nunchaku-lite-kernels), an unmodified
-  build of the same open-source kernels, and sets `DIFFUSERS_TRUST_REMOTE_KERNELS=true` unless you already set
-  it. Set `LOCAL_KERNELS` yourself to use a different build.
+- `transformer/modeling_nunchaku_qwenimage21.py`: the stock transformer class, unchanged.
+
+Both files start with the same kernel setup, so it runs whichever component loads first. Diffusers loads
+the NVFP4 kernels from `rootonchair/nunchaku-lite-kernels`, which is no longer downloadable. The setup
+points that name at [joseplcam/nunchaku-lite-kernels](https://huggingface.co/joseplcam/nunchaku-lite-kernels),
+an unmodified build of the same open-source kernels, and sets `DIFFUSERS_TRUST_REMOTE_KERNELS=true` unless
+you already set it. Set `LOCAL_KERNELS` yourself to use a different build.
 
 Use a different seed for an edit than the one that generated its input image. Qwen-Image-2.1 returns an
 over-sharpened copy that ignores the prompt when the edit starts from the same noise

text_encoder/modeling_nunchaku_qwen3vl.py

@@ -8,15 +8,31 @@ transformer. Everything else (vision tower, embeddings, the BF16 layers) loads
 as in the stock model.
 """
 
+import os
 from pathlib import Path
 
 import torch
 from accelerate import init_empty_weights
 from huggingface_hub import snapshot_download
 
-from .nunchaku_kernels import KERNELS  # noqa: F401  (sets up the kernels before Diffusers imports them)
-
-# isort: off  -- must stay below `.nunchaku_kernels`, which prepares the kernels this import loads
+# --- Kernel setup. Kept identical in transformer/modeling_nunchaku_qwenimage21.py: Diffusers downloads
+# only each component's own module file, and whichever component loads first must run it. ---
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True
+# --- end of kernel setup ---
+
+# isort: off  -- must stay below the kernel setup, which prepares the kernels this import loads
 from diffusers.quantizers.nunchaku.utils import replace_with_nunchaku_linear  # noqa: E402
 from safetensors.torch import load_file  # noqa: E402
 from transformers import Qwen3VLForConditionalGeneration  # noqa: E402

text_encoder/nunchaku_kernels.py

@@ -1,19 +0,0 @@
-# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
-#
-# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
-# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
-# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
-import os
-
-from huggingface_hub import snapshot_download
-
-KERNELS = "rootonchair/nunchaku-lite-kernels"
-if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
-    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
-    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
-if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
-    # Diffusers read this variable into a constant at import time, so set both.
-    import diffusers.utils.constants
-
-    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
-    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True

tools/modeling_nunchaku_qwen3vl.py

@@ -8,15 +8,31 @@ transformer. Everything else (vision tower, embeddings, the BF16 layers) loads
 as in the stock model.
 """
 
+import os
 from pathlib import Path
 
 import torch
 from accelerate import init_empty_weights
 from huggingface_hub import snapshot_download
 
-from .nunchaku_kernels import KERNELS  # noqa: F401  (sets up the kernels before Diffusers imports them)
-
-# isort: off  -- must stay below `.nunchaku_kernels`, which prepares the kernels this import loads
+# --- Kernel setup. Kept identical in transformer/modeling_nunchaku_qwenimage21.py: Diffusers downloads
+# only each component's own module file, and whichever component loads first must run it. ---
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True
+# --- end of kernel setup ---
+
+# isort: off  -- must stay below the kernel setup, which prepares the kernels this import loads
 from diffusers.quantizers.nunchaku.utils import replace_with_nunchaku_linear  # noqa: E402
 from safetensors.torch import load_file  # noqa: E402
 from transformers import Qwen3VLForConditionalGeneration  # noqa: E402

tools/modeling_nunchaku_qwenimage21.py

@@ -1,13 +1,32 @@
 """Qwen-Image-2.1 transformer, unchanged, loaded as a custom component.
 
-Its only job is importing `nunchaku_kernels`, so the NVFP4 kernels are set up even when
+Its only job is running the kernel setup below, so the NVFP4 kernels are ready even when
 Diffusers loads this component before the text encoder. The Nunchaku Lite quantization
 itself is handled by Diffusers' built-in quantizer from `config.json`.
 """
 
-from diffusers import QwenImage21Transformer2DModel
+import os
 
-from .nunchaku_kernels import KERNELS  # noqa: F401
+from huggingface_hub import snapshot_download
+
+# --- Kernel setup. Kept identical in text_encoder/modeling_nunchaku_qwen3vl.py: Diffusers downloads
+# only each component's own module file, and whichever component loads first must run it. ---
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True
+# --- end of kernel setup ---
+
+from diffusers import QwenImage21Transformer2DModel  # noqa: E402
 
 
 class NunchakuQwenImage21Transformer2DModel(QwenImage21Transformer2DModel):

tools/nunchaku_kernels.py

@@ -1,19 +0,0 @@
-# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
-#
-# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
-# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
-# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
-import os
-
-from huggingface_hub import snapshot_download
-
-KERNELS = "rootonchair/nunchaku-lite-kernels"
-if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
-    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
-    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
-if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
-    # Diffusers read this variable into a constant at import time, so set both.
-    import diffusers.utils.constants
-
-    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
-    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True

tools/package.py

@@ -27,14 +27,24 @@ from examples.convert_nunchaku_lite_diffusers import (  # noqa: E402
 )
 
 HERE = Path(__file__).parent
-# component -> (module file, class); both import nunchaku_kernels, which ships next to each
+# component -> (module file, class); both modules carry the same kernel setup block
 CUSTOM_COMPONENTS = {
     "text_encoder": ("modeling_nunchaku_qwen3vl", "NunchakuQwen3VLForConditionalGeneration"),
     "transformer": ("modeling_nunchaku_qwenimage21", "NunchakuQwenImage21Transformer2DModel"),
 }
 
 
+def kernel_setup_block(module: str) -> str:
+    text = (HERE / f"{module}.py").read_text()
+    body = text.split("# --- Kernel setup.", 1)[1].split("# --- end of kernel setup ---", 1)[0]
+    return body.split("---\n", 1)[1]  # drop the header line, which names the other file
+
+
 def main():
+    blocks = {kernel_setup_block(module) for module, _ in CUSTOM_COMPONENTS.values()}
+    if len(blocks) != 1:
+        raise ValueError("The kernel setup blocks in the two modeling files differ; keep them identical")
+
     parser = argparse.ArgumentParser()
     parser.add_argument("--dit", required=True)
     parser.add_argument("--text-encoder", required=True)
@@ -55,7 +65,6 @@ def main():
     index = json.loads((out / "model_index.json").read_text())
     for component, (module, cls) in CUSTOM_COMPONENTS.items():
         shutil.copy2(HERE / f"{module}.py", out / component / f"{module}.py")
-        shutil.copy2(HERE / "nunchaku_kernels.py", out / component / "nunchaku_kernels.py")
         index[component] = [module, cls]
     (out / "model_index.json").write_text(json.dumps(index, indent=2) + "\n")
 
@@ -66,7 +75,7 @@ def main():
     vae_config.pop("_name_or_path", None)  # a local path, meaningless on the Hub
     (out / "vae" / "config.json").write_text(json.dumps(vae_config, indent=2, sort_keys=True) + "\n")
 
-    for leftover in ("assets", ".gitattributes"):
+    for leftover in ("assets", ".gitattributes", ".cache"):  # .cache: snapshot_download bookkeeping
         path = out / leftover
         shutil.rmtree(path) if path.is_dir() else path.unlink(missing_ok=True)
     for doc in ("README.md", "NOTICE"):  # model card and license notice

transformer/modeling_nunchaku_qwenimage21.py

@@ -1,13 +1,32 @@
 """Qwen-Image-2.1 transformer, unchanged, loaded as a custom component.
 
-Its only job is importing `nunchaku_kernels`, so the NVFP4 kernels are set up even when
+Its only job is running the kernel setup below, so the NVFP4 kernels are ready even when
 Diffusers loads this component before the text encoder. The Nunchaku Lite quantization
 itself is handled by Diffusers' built-in quantizer from `config.json`.
 """
 
-from diffusers import QwenImage21Transformer2DModel
+import os
 
-from .nunchaku_kernels import KERNELS  # noqa: F401
+from huggingface_hub import snapshot_download
+
+# --- Kernel setup. Kept identical in text_encoder/modeling_nunchaku_qwen3vl.py: Diffusers downloads
+# only each component's own module file, and whichever component loads first must run it. ---
+# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
+# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
+# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
+KERNELS = "rootonchair/nunchaku-lite-kernels"
+if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
+    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
+    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
+if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
+    # Diffusers read this variable into a constant at import time, so set both.
+    import diffusers.utils.constants
+
+    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
+    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True
+# --- end of kernel setup ---
+
+from diffusers import QwenImage21Transformer2DModel  # noqa: E402
 
 
 class NunchakuQwenImage21Transformer2DModel(QwenImage21Transformer2DModel):

transformer/nunchaku_kernels.py

@@ -1,19 +0,0 @@
-# Shared by text_encoder/ and transformer/: whichever component Diffusers loads first sets this up.
-#
-# Diffusers fetches the Nunchaku Lite kernels from `rootonchair/nunchaku-lite-kernels`, which is
-# no longer downloadable. Loading this repo with trust_remote_code=True already runs this code,
-# so point that kernel name at the rebuilt copy before Diffusers imports its Nunchaku utilities.
-import os
-
-from huggingface_hub import snapshot_download
-
-KERNELS = "rootonchair/nunchaku-lite-kernels"
-if KERNELS not in os.environ.get("LOCAL_KERNELS", ""):
-    local = f"{KERNELS}={snapshot_download('joseplcam/nunchaku-lite-kernels')}"
-    os.environ["LOCAL_KERNELS"] = ":".join(filter(None, [os.environ.get("LOCAL_KERNELS"), local]))
-if "DIFFUSERS_TRUST_REMOTE_KERNELS" not in os.environ:
-    # Diffusers read this variable into a constant at import time, so set both.
-    import diffusers.utils.constants
-
-    os.environ["DIFFUSERS_TRUST_REMOTE_KERNELS"] = "true"
-    diffusers.utils.constants.DIFFUSERS_TRUST_REMOTE_KERNELS = True