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.
4. mingyi456 commented on Oct 24, 2026

   

   mingyi456

   on Oct 24, 2026

   Last edited by mingyi456

   ContributorAuthor

   More actions

   @Jackmeson1 In terms of code complexity and the overall motivation, I think it does make sense to do this, but that would constitute a change in the scope of my PR, which was originally meant for CosmosTransformer3DModel only.

   Expanding the scope of an in-progress PR would mean it needs to be reviewed again (i.e. if my PR were to be reviewed and approved as a change specifically for CosmosTransformer3DModel, it does not make sense for me to extend the changes for QwenImageTransformer2DModel and others between the approval and merging, even if the nature of the change is identical), but anyway I am not sure about the review status of my PR currently. @DN6 Can you weigh in on this?
5. Jackmeson1 commented on Oct 24, 2026

   

   Jackmeson1

   on Oct 24, 2026

   More actions

   Thanks, agreed: #14931 should stay Cosmos-only. I meant a separate small PR for the Qwen and MotifVideo mappings, once a maintainer confirms they want it.
6. mingyi456 commented on Oct 25, 2026

   

   mingyi456

   on Oct 25, 2026

   ContributorAuthor

   More actions

   No, I am saying I am not sure about the review status of my PR. If you can take a look, there is some maintainer activity but no explicit approval / requested changes, even though the scope of the PR is very straightforward.

   So I am asking if the maintainer has any opinions on expanding the PR scope to cover more models, I am not answering with a hard "yes" or "no".
7. Jackmeson1 commented on Oct 25, 2026

   

   Jackmeson1

   on Oct 25, 2026

   More actions

   Understood, thanks for clarifying. I'll wait for a maintainer's view on the scope before opening anything.