How to load the structured-pruned sd3-m using SD3Transformer2DModel.from_pretrained?

[plustwo0]

Hi, thanks for your great work!

I've tried your unstructured pruning and it turns out really efficient.
However, performing the structured one should reduce param. But when I load it with code below, the memory usage after loading the model remained the same as the original unpruned version (using the code below and without ignore_mismatched_sizes=True it would report a bug)

Is there a way to load the pruned model so that it actually uses less memory?

transformer = SD3Transformer2DModel.from_pretrained(
            args.model_path,
            ignore_mismatched_sizes=True,
            low_cpu_mem_usage=False
            )

[Alrightlone]

Hi, thanks for your great work!

I've tried your unstructured pruning and it turns out really efficient. However, performing the structured one should reduce param. But when I load it with code below, the memory usage after loading the model remained the same as the original unpruned version (using the code below and without ignore_mismatched_sizes=True it would report a bug)

Is there a way to load the pruned model so that it actually uses less memory?

transformer = SD3Transformer2DModel.from_pretrained(
            args.model_path,
            ignore_mismatched_sizes=True,
            low_cpu_mem_usage=False
            )

Hi, thanks for your feedback!

Since structured pruning physically changes the model architecture (e.g., layer dimensions), loading it via from_pretrained relies on the original model configuration, which will attempt to initialize the model with the original shapes. This is why you see no memory reduction and need ignore_mismatched_sizes=True.

To load the structurally pruned model and benefit from the memory reduction, you should use torch.load directly, as the modified model structure is saved as a complete object in main.py (torch.save(pipe.transformer, args.save_model)).

You can load the pruned checkpoint for inference like this:

import torch

# ... existing initialization code ...

# Load the pruned transformer structure and weights directly
pipe.transformer = torch.load("/path/to/your/checkpoint", weights_only=False)

# Now the pipe uses the pruned model with reduced memory footprint

Hope this helps!