error when using _init_tiled_vae in pasd_sdxl

[1vanFan]

After adding the VAEHook in test_pasd_sdxl, I encountered an error when running it.

TypeError: Linear.forward() got an unexpected keyword argument 'scale'

python test_pasd_sdxl.py -i datasets/test -o output/test --num_inference_steps 2 -s 2 --decoder_tiled_size 128

INFO:root:Loaded coca_ViT-L-14 model config.
INFO:root:Loading pretrained coca_ViT-L-14 weights (mscoco_finetuned_laion2B-s13B-b90k).
1 / 1 datasets/test/662.png a woman sitting on a couch with her hands folded on her chest . photorealistic, clean, high-resolution, 8k
100%|████████████████████████████████████████████████| 2/2 [00:09<00:00,  4.86s/it]
[Tiled VAE]: input_size: torch.Size([1, 4, 169, 165]), tile_size: 128, padding: 11
[Tiled VAE]: split to 2x2 = 4 tiles. Optimal tile size 96x96, original tile size 128x128
[Tiled VAE]: Executing Decoder Task Queue:   0%|           | 0/492 [00:00<?, ?it/s]Traceback (most recent call last):
  File "/data/bear/PASD/test_pasd_sdxl.py", line 381, in <module>
    main(args)
  File "/data/bear/PASD/test_pasd_sdxl.py", line 281, in main
    image = pipeline(
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/bear/PASD/pipelines/pipeline_pasd_sdxl.py", line 1639, in __call__
    image = self.vae.decode(latents, return_dict=False)[0]
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py", line 46, in wrapper
    return method(self, *args, **kwargs)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 321, in decode
    decoded = self._decode(z).sample
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 292, in _decode
    dec = self.decoder(z)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/bear/PASD/myutils/vaehook.py", line 691, in __call__
    return self.vae_tile_forward(x)
  File "/data/bear/PASD/myutils/vaehook.py", line 576, in wrapper
    ret = fn(*args, **kwargs)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/bear/PASD/myutils/vaehook.py", line 927, in vae_tile_forward
    tile = task[1](tile)
  File "/data/bear/PASD/myutils/vaehook.py", line 363, in <lambda>
    task_queue.append(('attn', lambda x, net=net: attn_forward_new_xformers(net, x)))
  File "/data/bear/PASD/myutils/vaehook.py", line 266, in attn_forward_new_xformers
    query = self.to_q(hidden_states, scale=scale)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/adminroot/anaconda3/envs/bear_supir/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
TypeError: Linear.forward() got an unexpected keyword argument 'scale'
[Tiled VAE]: Executing Decoder Task Queue:   8%| | 41/492 [00:00<00:02, 199.17it/s]

I tried commenting out the scale parameter in the query = self.to_q(hidden_states, scale=scale) method. But the images show noise.

[SafetyLucas]

The scale argument problem can be fixed by removing any mention of scale in vaehook.py. The artifact problem only occurs when tiled vae is used. I think the repo he copied code for vaehook from was not well implemented for SDXL. I saw several other posts where problems were caused by invisible watermark, but those looked more like red dots/splotches. The noise seen in your image (and mine as well) appear to be a vae problem.