ostris

/

Flex.2-preview

like 206

This used to work with the alpha, but it doesn't work with the Preview:

import torch
#image_model_card = "ostris/Flex.1-alpha"
image_model_card = "ostris/Flex.2-preview"

from diffusers import BitsAndBytesConfig, FluxTransformer2DModel, FluxPipeline

nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model_nf4 = FluxTransformer2DModel.from_pretrained(
    image_model_card,
    subfolder="transformer",
    quantization_config=nf4_config,
    torch_dtype=torch.bfloat16,
)
pipe = FluxPipeline.from_pretrained(
    image_model_card,
    transformer=model_nf4,
    torch_dtype=torch.bfloat16,
)

pipe.enable_model_cpu_offload()

image = pipe(
    prompt="a cat with a hat",
    prompt_2=None,
    max_sequence_length=512,
    num_inference_steps=25,
    guidance_scale=3,
    height=720,
    width=1280,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save("flex_2_preview.png")

This is the error introduced with the Preview:

Error: Python: Traceback (most recent call last):
  File "...\python\Lib\site-packages\torch\utils\_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "...\python\Lib\site-packages\diffusers\pipelines\flux\pipeline_flux.py", line 942, in __call__
    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "...\python\Lib\site-packages\diffusers\schedulers\scheduling_flow_match_euler_discrete.py", line 461, in step
    prev_sample = sample + dt * model_output
                  ~~~~~~~^~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (196) must match the size of tensor b (64) at non-singleton dimension 2