Generate Different Number of Images Without Recompilation

[liorRabkin]

I'm trying to generate a different number of images without recompilation. For example, warm up with 1 image and then generate with 2 images. I tried with dynamic=True, but I still get recompilation when changing num_images_per_prompt:

with torch._dynamo.config.patch(error_on_recompile=True):
      pipe = DiffusionPipeline.from_pretrained(
          "black-forest-labs/FLUX.1-dev,
          torch_dtype=torch.bfloat16,
      ).to("cuda")

      pipe.transformer = torch.compile(pipe.transformer, dynamic=True)

      # warmup
      image = pipe(prompt="A cat holding a sign that says hello world",
                   width=1024,
                   height=1024,
                   num_inference_steps=10,
                   num_images_per_prompt=1,
                   ).images[0]

      # inference
      image = pipe(prompt="A cat holding a sign that says hello world",
                   width=1024,
                   height=1024,
                   num_inference_steps=10,
                   num_images_per_prompt=2,
                   ).images[0]

The error:

torch._dynamo.exc.RecompileError: Recompiling function forward in /root/PycharmProjects/project/.venv/lib/python3.10/site-packages/diffusers/models/transformers/transformer_flux.py:389
    triggered by the following guard failure(s):
    - 0/0: tensor 'guidance' size mismatch at index 0. expected 1, actual 2

[sayakpaul]

Cc: @anijain2305

[anijain2305]

It sounds like a 0/1 specialization to me. @laithsakka Is there a way to avoid 0/1 specialization?

cc @StrongerXi

[sayakpaul]

Was there any update on this?

[laithsakka]

@liorRabkin

1. can you try mark_unbacked? marking things dynamic does not guarantee no recompilations. while mark_unbacked does guarantee no guards are added hence no recompilation. (might hit DDE in user code and hopefully wont?). we are almost done making framework code more unbacked friendly. can you give it a try.

2. less reliable can you try to use the option backed_size_oblivious? we still have to harden it and might deprecate but worth a try

[StrongerXi]

@liorRabkin also note that the recompilation will only happen once for the initial num_images_per_prompt change from 1.