CosmosTransformer3DModel.from_single_file() does not work with the model.diffusion_model. prefix #14929

Activity

1. mingyi456

   added

   bugSomething isn't workingSomething isn't working

   on Oct 3, 2026
2. github-actions

   added

   models

   single-file

   on Oct 3, 2026
3. Jackmeson1 commented on Oct 20, 2026

   

   Jackmeson1

   on Oct 20, 2026

   More actions

   I reproduced an additional model.diffusion_model. loading gap in Qwen and MotifVideo that may belong alongside #14931.

   This is a synthetic CPU loader regression case with tiny random models. I have not rerun the full Comfy-Org fp8 checkpoints here.

   Environment: Diffusers e04f0c5f / 0.42.0.dev0, Python 3.12.15, PyTorch 2.14.1+cpu, Transformers 5.20.1, Accelerate 1.15.0, Safetensors 0.8.0, Linux/WSL2. No downloads or GPU are needed:

   import tempfile
   from pathlib import Path
   
   import torch
   from safetensors.torch import save_file
   from diffusers import QwenImageTransformer2DModel, MotifVideoTransformer3DModel
   from diffusers.utils import logging
   
   logging.set_verbosity_error()
   torch.manual_seed(0)
   models = [
       QwenImageTransformer2DModel(
           patch_size=2, in_channels=16, out_channels=4, num_layers=2,
           attention_head_dim=8, num_attention_heads=2,
           joint_attention_dim=16, axes_dims_rope=(2, 2, 4),
       ),
       MotifVideoTransformer3DModel(
           in_channels=33, out_channels=16, num_attention_heads=2,
           attention_head_dim=12, num_layers=1, num_single_layers=1,
           num_decoder_layers=0, text_embed_dim=32, image_embed_dim=4,
           rope_axes_dim=(4, 4, 4),
           enable_text_cross_attention_dual=False,
           enable_text_cross_attention_single=False,
       ),
   ]
   for source in models:
       cls = type(source)
       state = {k: v.contiguous() for k, v in source.state_dict().items()}
       with tempfile.TemporaryDirectory() as directory:
           source.save_config(directory)
           for label, prefix in [("plain", ""), ("prefixed", "model.diffusion_model.")]:
               checkpoint = {prefix + k: v for k, v in state.items()}
               path = str(Path(directory) / f"{label}.safetensors")
               save_file(checkpoint, path)
               try:
                   loaded = cls.from_single_file(path, config=directory, local_files_only=True)
                   actual = loaded.state_dict()
                   meta = sum(v.is_meta for v in actual.values())
                   equal = not meta and actual.keys() == state.keys() and all(
                       torch.equal(actual[k], v) for k, v in state.items()
                   )
                   print(cls.__name__, label, f"meta={meta}", f"equal={equal}")
               except Exception as error:
                   print(cls.__name__, label, type(error).__name__, str(error).split(";")[0])
           stripped = {k.removeprefix("model.diffusion_model."): v for k, v in checkpoint.items()}
           loaded = cls.from_single_file(stripped, config=directory, local_files_only=True)
           assert all(torch.equal(loaded.state_dict()[k], v) for k, v in state.items())
           print(cls.__name__, "stripped: all tensors equal")

   Observed:

   Class	Plain	Prefixed	Manually stripped
   QwenImageTransformer2DModel	meta=0 equal=True	NotImplementedError: Cannot copy out of meta tensor	All tensors equal
   MotifVideoTransformer3DModel	meta=0 equal=True	Same error	All tensors equal

   Expected: adding this checkpoint namespace should preserve the loaded weights, as supported by several existing transformer converters.

   Both classes use identity checkpoint_mapping_fn entries, so prefixed keys remain unmatched and parameters stay on meta until dispatch fails. QwenImage21 already removes this prefix.

   Would it make sense to cover these mappings alongside #14931? I can send a PR after maintainer acknowledgment.

   Found and reproduced with AI assistance; the repro above was run by me.