Unet sample_size not fully considered in pipeline init

[Foundsheep]

Describe the bug

Background

• I'm trying to train a small unet from scratch
• For this, I used a modified version of train_text_to_image.py
• unet model's sample_size is not symmetric. I passed [60, 80] when instantiating UNet2DConditionModel, which is normally an int value with 512

What happens

• An error occurs in the validation step, using log_validation()
• It's because the pipeline(here, StableDiffusionPipeline) loaded in the method doesn't expect unet.config.sample_size to be a list, but an int when checking for deprecation_message.
• But I think it should be able to accept list which is already one of the properly accepted ways to pass a value for UNet2DConditionModel's sample_size parameter.
• The exact line where the error occurs is the code below

# in `pipeline_stable_diffusion.py` line 258
is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64

Other things to mention

• It seems like that line of code is spread across many other files, such as other pipeline_stable...py.
• If you want, I can spend some time on this and make a PR
• If so, I'm thinking of making that line of code with other relevant lines into a function and put that function somewhere such as in utils(the location needs to be specified by other members of HF staff)

Reproduction

from diffusers import StableDiffusionPipeline, UNet2DConditionModel

# load pre-trained model
sd_repo = "stable-diffusion-v1-5/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(sd_repo, use_safetensors=True)

# prepare components
small_unet = UNet2DConditionModel(sample_size=[60, 80], block_out_channels=[32, 64, 128, 256])
vae = pipe.components["vae"]
text_encoder = pipe.components["text_encoder"]
tokenizer = pipe.components["tokenizer"]
scheduler = pipe.components["scheduler"]
safety_checker = None
feature_extractor = pipe.components["feature_extractor"]
image_encoder = pipe.components["image_encoder"]

# error occurs(this would be the same when loading with .from_pretrained())
new_pipe = StableDiffusionPipeline(
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    unet=small_unet,
    scheduler=scheduler,
    safety_checker=safety_checker,
    feature_extractor=feature_extractor,
    image_encoder=image_encoder
)

Logs

[rank0]: Traceback (most recent call last):
[rank0]:   File "/root/workspace/03_gen/generative_02/diffusers/examples/text_to_image/train_text_to_image_with_new_unet_from_scratch.py", line 1208, in <module>
[rank0]:     main()
[rank0]:   File "/root/workspace/03_gen/generative_02/diffusers/examples/text_to_image/train_text_to_image_with_new_unet_from_scratch.py", line 1143, in main
[rank0]:     log_validation(
[rank0]:   File "/root/workspace/03_gen/generative_02/diffusers/examples/text_to_image/train_text_to_image_with_new_unet_from_scratch.py", line 144, in log_validation
[rank0]:     pipeline = StableDiffusionPipeline.from_pretrained(
[rank0]:   File "/root/workspace/03_gen/generative_02/.venv/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 114, in _inner_fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/root/workspace/03_gen/generative_02/.venv/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py", line 948, in from_pretrained
[rank0]:     model = pipeline_class(**init_kwargs)
[rank0]:   File "/root/workspace/03_gen/generative_02/.venv/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 258, in __init__
[rank0]:     is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
[rank0]: TypeError: '<' not supported between instances of 'list' and 'int'

### System Info

- 🤗 Diffusers version: 0.32.0.dev0
- Platform: Linux-5.4.0-167-generic-x86_64-with-glibc2.31
- Running on Google Colab?: No
- Python version: 3.10.10
- PyTorch version (GPU?): 2.5.1+cu124 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Huggingface_hub version: 0.26.2
- Transformers version: 4.46.3
- Accelerate version: 1.1.1
- PEFT version: 0.7.0
- Bitsandbytes version: not installed
- Safetensors version: 0.4.5
- xFormers version: not installed
- Accelerator: GRID A100X-10C, 10240 MiB
GRID A100X-10C, 10240 MiB
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@yiyixuxu 

[hlky]

Hi @Foundsheep! Thanks for finding this. For context this check was added ~2 years ago due to some incorrect configs so I don't think we need to apply the check for tuple case, we can add isinstance(unet.config.sample_size, int):

is_unet_sample_size_less_64 = (
    hasattr(unet.config, "sample_size")
    and isinstance(unet.config.sample_size, int)
    and unet.config.sample_size < 64
)

Would you like to open a PR?

[Foundsheep]

@hlky Thanks for the quick reply.

Sure, :) I'd like to contribute. I'll open a PR, but just to make sure, it seems like that code is in 70 places. Do you think it is better for me to replace all those occurrences?

[hlky]

First let's check if any other changes are required, can you make the change to StableDiffusionPipeline and test it? Looks like this will need changing

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

Lines 905 to 906 in 89e4d62

	height = height or self.unet.config.sample_size * self.vae_scale_factor
	width = width or self.unet.config.sample_size * self.vae_scale_factor

Also in UNet2DConditionModel there's a discrepancy between docstring and type hint

diffusers/src/diffusers/models/unets/unet_2d_condition.py

Lines 82 to 83 in 6131a93

	sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
	Height and width of input/output sample.

diffusers/src/diffusers/models/unets/unet_2d_condition.py

Line 173 in 6131a93

sample_size: Optional[int] = None,

[Foundsheep]

Ok, I've made some changes where needed, and they worked ok when I tested them in my local environment.
Here are the changes

        self._is_unet_config_sample_size_int = isinstance(unet.config.sample_size, int)
        is_unet_sample_size_less_64 = (
            hasattr(unet.config, "sample_size") 
            and self._is_unet_config_sample_size_int 
            and unet.config.sample_size < 64
        )

https://github.com/Foundsheep/diffusers/blob/36fb12b9682e6653dc51f5011f70d307215f2f4a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L258-L263

        # 0. Default height and width to unet
        if not height or not width:
            height = self.unet.config.sample_size if self._is_unet_config_sample_size_int else self.unet.config.sample_size[0]
            width = self.unet.config.sample_size if self._is_unet_config_sample_size_int else self.unet.config.sample_size[1]
            height, width = height * self.vae_scale_factor, width * self.vae_scale_factor
        # to deal with lora scaling and other possible forward hooks

https://github.com/Foundsheep/diffusers/blob/36fb12b9682e6653dc51f5011f70d307215f2f4a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L909-L914

        sample_size: Optional[Union[int, Tuple[int, int]]] = None,

https://github.com/Foundsheep/diffusers/blob/17c7889479407df52f129552469b58d8b3653693/src/diffusers/models/unets/unet_2d_condition.py#L173

[hlky]

Thanks! We can create a PR for these changes to StableDiffusionPipeline and UNet2DConditionModel then replicate to other locations after further review

[Foundsheep]

@hlky Thanks for the guidance.

As I mentioned earlier, the same line of code spreads across the project in about 70 places and just one of them has been modified by the recent merge. Do you think the same modification needs to be done in all the cases now?