ValueError: Invalid sample size issue when testing 2K checkpoint

[MahdiMohseni0033]

Thanks for your amazing work, I have faced an error during testing the 2K checkpoint :

ValueError: Invalid sample size 

I wish you help me with this error, waiting for your response

[lawrence-cj]

What's your test script?

[MahdiMohseni0033]

following is my code, similar to one suggested in Readme:

import torch
from diffusers import Transformer2DModel
from scripts.diffusers_patches import pixart_sigma_init_patched_inputs, PixArtSigmaPipeline
import time 
assert getattr(Transformer2DModel, '_init_patched_inputs', False), "Need to Upgrade diffusers: pip install git+https://github.com/huggingface/diffusers"
setattr(Transformer2DModel, '_init_patched_inputs', pixart_sigma_init_patched_inputs)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weight_dtype = torch.float16

transformer = Transformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-2K-MS", 
    subfolder='transformer', 
    torch_dtype=weight_dtype,
    use_safetensors=True,
)
pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
    transformer=transformer,
    torch_dtype=weight_dtype,
    use_safetensors=True,
)
pipe.to(device)

prompt = "a dog walking in the beach"
tik = time.time()
image = pipe(prompt).images[0]
print(f'Generation time :{time.time() - tik}')
image.save(f"result.png")

[MahdiMohseni0033]

it solved by using diffusers version :

from diffusers import PixArtAlphaPipeline, PixArtSigmaPipeline, Transformer2DModel, DDPMScheduler
import torch

transformer = Transformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-2K-MS",
    subfolder="transformer",
    torch_dtype=torch.float16,
    use_additional_conditions=False,
)
pipe = PixArtSigmaPipeline.from_pretrained("PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers", transformer=transformer, torch_dtype=torch.float16).to('cuda')
pipe.to('cuda')
# pipe.enable_model_cpu_offload()
prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe(prompt).images[0]
image.save("yiyi_test_4_out_sigma_2k.png")

[lawrence-cj]

Cool cool. We just got merged into diffusers yesterday. It should work. Thx anyway.