Compilation fails after fb84060

[kasper93]

After changes in fb84060 compilation fails, works fine before.

python unet.py --hf-hub-or-path "runwayml/stable-diffusion-v1-5" --width 512 512 --height 512 512 --batch-size 1 1 --clip-chunks 10 --model-name "v15_unet_512_512" --work-dir "/app/tmp/"

Traceback (most recent call last):
  File "AITemplate/unet.py", line 141, in <module>
    compile_diffusers()
  File "venv/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "AITemplate/unet.py", line 104, in compile_diffusers
    compile_unet(
  File "AITemplate/ait/compile/unet.py", line 271, in compile_unet
    Y = ait_mod(
        ^^^^^^^^
  File "venv/lib/python3.11/site-packages/aitemplate/frontend/nn/module.py", line 389, in _call_impl
    return forward_call(*input, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "AITemplate/ait/modeling/unet_2d_condition.py", line 394, in forward
    sample = upsample_block(
             ^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/aitemplate/frontend/nn/module.py", line 389, in _call_impl
    return forward_call(*input, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "AITemplate/ait/modeling/unet_blocks.py", line 671, in forward
    hidden_states = upsampler(hidden_states, upsample_size)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/aitemplate/frontend/nn/module.py", line 389, in _call_impl
    return forward_call(*input, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "AITemplate/ait/modeling/resnet.py", line 71, in forward
    x = nn.Upsampling2d(scale_factor=2.0, mode="nearest")(x, out)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/aitemplate/frontend/nn/module.py", line 389, in _call_impl
    return forward_call(*input, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "venv/lib/python3.11/site-packages/aitemplate/frontend/nn/upsample.py", line 47, in forward
    assert len(args) == 1
           ^^^^^^^^^^^^^^
AssertionError

[finlay96]

Is there any update on this at all as problem is still persisting?

[IdiotSandwichTheThird]

Sure would be nice if this were fixed, considering there's no working modules > batch size 1 for 1.5.