Support DiffusionPipeline.from_pretrained(repo, trust_remote_code=True)

model_index.json uses the [file, class] form, pipeline.py imports transformer.py relatively so the
dynamic-module loader fetches it too, and the transformer class is published on the diffusers module
(the stock component loader has no auto_map/trust_remote_code path). Tokenizer resolved from the
processor before register_modules, VAE pinned to fp32 on every entry point.

README.md

@@ -74,10 +74,11 @@ Every image below is generated by this pipeline with 30 steps at 1024 px.
 
 ```python
 import torch
-from pipeline import ZenImageEditPipeline      # shipped in this repo
+from diffusers import DiffusionPipeline
 
-pipe = ZenImageEditPipeline.from_pretrained(".", dtype=torch.float16)
-pipe.enable_model_cpu_offload()                # 14.5 GB DiT + fp32 VAE decoder do not co-reside on 32 GB
+pipe = DiffusionPipeline.from_pretrained("AiArtLab/zen-image-edit", trust_remote_code=True,
+                                         dtype=torch.float16)
+pipe.enable_model_cpu_offload()   # 14.5 GB DiT + fp32 VAE decoder do not co-reside on 32 GB
 
 # text-to-image
 image = pipe(prompt="a red fox in a snowy forest at dusk, cinematic, 85mm",
@@ -92,8 +93,21 @@ image = pipe(prompt="Replace the woman in <image2> with the woman from <image1>;
              generator=torch.Generator("cuda").manual_seed(1234)).images[0]
 ```
 
+`trust_remote_code=True` pulls `pipeline.py` and `transformer.py` from this repo and runs them, so
+no clone is needed. Cloning works too and gives the class directly:
+
+```python
+from pipeline import ZenImageEditPipeline
+pipe = ZenImageEditPipeline.from_pretrained(".", dtype=torch.float16)
+```
+
 CLI: `python example.py --prompt "..." [--image a.png b.png] --out out.png`
 
+Requirements: `torch`, `transformers`, `accelerate` and a `diffusers` built with Qwen-Image-2.1
+(`pip install git+https://github.com/huggingface/diffusers`) — the transformer subclasses
+`QwenImage21Transformer2DModel`. `trust_remote_code` saves the clone, it does **not** save the 17 GB
+of weights.
+
 ### Files
 
 ```
@@ -110,8 +124,9 @@ media/               the examples above
 ```
 
 `QwenImage21FusionTransformer2DModel` is a custom class defined in `transformer.py`, not registered
-inside `diffusers`, so plain `DiffusionPipeline.from_pretrained` does not resolve it. Load through the
-shipped pipeline with this folder on `sys.path`.
+inside `diffusers`, so the pipeline publishes it on the `diffusers` module at import time. That is
+what makes the `trust_remote_code=True` one-liner above work; without it the stock component loader
+would not find the DiT class.
 
 ### Limitations
 

model_index.json

@@ -1,5 +1,8 @@
 {
-  "_class_name": "ZenImageEditPipeline",
+  "_class_name": [
+    "pipeline",
+    "ZenImageEditPipeline"
+  ],
   "_diffusers_version": "0.41.0.dev0",
   "processor": [
     "transformers",
@@ -13,10 +16,6 @@
     "transformers",
     "Qwen3_5ForConditionalGeneration"
   ],
-  "tokenizer": [
-    "transformers",
-    "Qwen2Tokenizer"
-  ],
   "transformer": [
     "diffusers",
     "QwenImage21FusionTransformer2DModel"

pipeline.py

@@ -16,7 +16,22 @@ from PIL import Image as PILImage
 from diffusers.image_processor import VaeImageProcessor
 from diffusers.pipelines.qwenimage21.pipeline_qwenimage21 import QwenImage21Pipeline
 
-from transformer import QwenImage21FusionTransformer2DModel
+try:
+    # Loaded as a Hub dynamic module (`trust_remote_code=True`): the loader follows relative imports
+    # and fetches `transformer.py` from the same repo next to this file.
+    from .transformer import QwenImage21FusionTransformer2DModel
+except ImportError:
+    # Plain module on `sys.path` (local use, `example.py`).
+    from transformer import QwenImage21FusionTransformer2DModel
+
+# `DiffusionPipeline.from_pretrained(repo, trust_remote_code=True)` resolves every component by
+# looking its class up in the library named by `model_index.json`, and diffusers has no
+# `auto_map`/`trust_remote_code` path for model components (`models/model_loading_utils.py`).
+# Publishing the class on the `diffusers` module makes the stock loader find it: this module is
+# imported while the pipeline class is resolved, i.e. before any component is loaded.
+import diffusers  # noqa: E402
+
+diffusers.QwenImage21FusionTransformer2DModel = QwenImage21FusionTransformer2DModel
 
 
 class ZenImageEditPipeline(QwenImage21Pipeline):
@@ -44,6 +59,18 @@ class ZenImageEditPipeline(QwenImage21Pipeline):
         # `processor.apply_chat_template(system)`, and the Qwen3.5 processor raises
         # "No user query found" on a system-only message. The rest is its own code.
         super(QwenImage21Pipeline, self).__init__()
+        if isinstance(tokenizer, (list, tuple)):
+            # Loading through `DiffusionPipeline.from_pretrained` hands tokenizer components over as
+            # a `[slow, fast]` pair; we only need one.
+            tokenizer = next((t for t in tokenizer if t is not None), None)
+        if tokenizer is None:
+            # The Hub loader does not pass a tokenizer (it lives inside the processor); resolve it
+            # before `register_modules`, otherwise `self.tokenizer` would silently stay None.
+            tokenizer = getattr(processor, "tokenizer", None)
+        if vae is not None:
+            # The repo ships the VAE in fp32, and loading through the Hub applies one `torch_dtype`
+            # to every component. Pin it here so both entry points behave identically.
+            vae.to(torch.float32)
         self.register_modules(
             vae=vae,
             text_encoder=text_encoder,
@@ -75,7 +102,6 @@ class ZenImageEditPipeline(QwenImage21Pipeline):
         if not self.student_layers:
             raise ValueError("transformer config has no text_fusion_config.student_layers")
         self._drop_idx = int(fusion.get("drop_idx", 0))
-        tokenizer = tokenizer if tokenizer is not None else getattr(processor, "tokenizer", None)
         self._img_token_id = int(tokenizer.encode("<|image_pad|>", add_special_tokens=False)[0])
         # The prefix crop happens inside the transformer; if the tokenizer and the config ever
         # disagree the condition silently slides by a few positions. Check once, at build time.
@@ -162,12 +188,16 @@ class ZenImageEditPipeline(QwenImage21Pipeline):
 
     @classmethod
     def from_pretrained(cls, root=".", dtype=torch.float16, **kwargs):
-        """Load the pipeline from a model folder (the layout shipped in this repo)."""
+        """Load the pipeline from a local model folder (the layout shipped in this repo).
+
+        For a Hub repo use `DiffusionPipeline.from_pretrained(repo_id, trust_remote_code=True)`,
+        which resolves this class through `model_index.json` and loads the components itself.
+        """
         from diffusers import AutoencoderKLQwenImage21, FlowMatchEulerDiscreteScheduler
         from transformers import AutoProcessor, AutoTokenizer, Qwen3_5ForConditionalGeneration
 
         if not os.path.isdir(root):
-            raise ValueError(f"expected a model folder with the components, got {root!r}")
+            raise ValueError(f"expected a local model folder with the components, got {root!r}")
         transformer = QwenImage21FusionTransformer2DModel.from_pretrained(
             os.path.join(root, "transformer"), torch_dtype=dtype
         )