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

[oghetti]

Error Message:
Traceback (most recent call last): File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/processors/actor_base.py", line 139, in process_messages out_payload_or_list = self.process_payload(payload) # <- abstract method File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/common/logging.py", line 104, in _wrap return f(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/processors/enhancer.py", line 59, in process_payload out_images: list[Image] = [ File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/processors/enhancer.py", line 60, in <listcomp> self._process_image(image, model_input, pre=f"Image [{i}/{len(in_images)}]") File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/processors/enhancer.py", line 71, in _process_image out_pil: PILImage = self._model.process(in_pil, model_input) File "/app/.venv/lib/python3.10/site-packages/aimc_toolkit/common/logging.py", line 104, in _wrap return f(*args, **kwargs) File "/app/aimc_model_pasd/model.py", line 146, in process image_enh: PILImage = self._model( File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 581, in wrapper ret = fn(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/pasd/pipelines/pipeline_pasd.py", line 1188, in __call__ image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]#.flip(1) File "/app/.venv/lib/python3.10/site-packages/diffusers/utils/accelerate_utils.py", line 46, in wrapper return method(self, *args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 314, in decode decoded = self._decode(z).sample File "/app/.venv/lib/python3.10/site-packages/diffusers/models/autoencoders/autoencoder_kl.py", line 285, in _decode dec = self.decoder(z) File "/app/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 698, in __call__ return self.vae_tile_forward(x) File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 581, in wrapper ret = fn(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context return func(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 931, in vae_tile_forward tile = task[1](tile) File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 367, in <lambda> task_queue.append(('attn', lambda x, net=net: attn_forward_new_xformers(net, x))) File "/app/.venv/lib/python3.10/site-packages/pasd/myutils/vaehook.py", line 264, in attn_forward_new_xformers query = self.to_q(hidden_states, scale=scale) File "/app/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl return self._call_impl(*args, **kwargs) File "/app/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl return forward_call(*args, **kwargs) TypeError: Linear.forward() got an unexpected keyword argument 'scale'

Environment to reproduce:

Python==3.10.13
diffusers==0.29.2
torch==2.2.0
CUDA Version 11.8

In the forementioned method, scale is set by default to 1 and it's never overwritten. Plus, the scaled dot product attention is also calculated without setting any explicit scale.
hidden_states = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False )
Given that, do you think it is possibile to remove reference of scale when building q, k and v Tensors?

[dezson]

Hey, is there any update on this? 🤔

[oghetti]

@yangxy @Atry A PR is ready to solve this issue, see: PASD

Can we merge this to main?

[hazebroufe]

I had the same problem. Just removing the scale arguments in vaehook.py (from where the problematic forward is called) worked for me.
Here is the bash script that creates the corresponding python file and runs it to remove the scale parameters (just copy-past the full block into terminal from the PASD main folder):

# Create a fix script
cat > fix_vaehook_scale.py << 'EOF'
import re

file_path = "pasd/models/pasd/vaehook.py"

with open(file_path, 'r') as f:
    content = f.read()

# Remove scale parameter from all Linear layer calls
# Pattern: matches calls like .to_q(x, scale=scale) and replaces with .to_q(x)
patterns = [
    (r'self\.to_q\(([^,]+), scale=scale\)', r'self.to_q(\1)'),
    (r'self\.to_k\(([^,]+), scale=scale\)', r'self.to_k(\1)'),
    (r'self\.to_v\(([^,]+), scale=scale\)', r'self.to_v(\1)'),
    (r'self\.to_out\[0\]\(([^,]+), scale=scale\)', r'self.to_out[0](\1)'),
]

for pattern, replacement in patterns:
    content = re.sub(pattern, replacement, content)

with open(file_path, 'w') as f:
    f.write(content)

print("✓ Fixed vaehook.py - removed scale parameters from Linear layer calls")
EOF

python fix_vaehook_scale.py