Fix for AutoModel.from_pretrained crashing

[EviLooo]

Title

Fix for AutoModel.from_pretrained crashing on Google Colab (trust_remote_code, dtype, and missing attention fallback)

Description

If you are trying to run AutoModel.from_pretrained("ShandaAI/FloodDiffusionTiny") in a Google Colab notebook or a fresh environment, it currently crashes due to four issues with dependencies and the cached HuggingFace files:

1. Incompatible Transformers Version: Modern versions of transformers (v5.x) crash when trying to load google/umt5-base. You must downgrade to a specific v4 version (like 4.44.0) that supports umt5 without breaking tokenizers on Python 3.12.
2. trust_remote_code check fails: hf_pipeline.py checks kwargs.get('trust_remote_code'), but transformers strips this out before passing it to the model.
3. dtype TypeError: diffusion_forcing_wan_tiny.py passes dtype= to the text encoder, but newer versions of transformers strictly require torch_dtype=.
4. Missing Native Attention Fallback: tools/attention.py forces the use of Flash Attention. If a user doesn't have it (for example collab/kaggle T4 gpu doesn't support it), the code just hits assert FLASH_ATTN_2_AVAILABLE and crashes. Furthermore, writing a custom PyTorch fallback fails because the variable lk defaults to q_lens (length 65), even though the text key tensor k is only length 1 for cross-attention, causing severe shape mismatches.

The Fix

Step 1: Force install a compatible transformers version in your terminal/notebook. (Note: You can optionally add flash-attn --no-build-isolation to this command if your GPU supports it).

pip install transformers==4.44.0

Step 2: Run this Python script immediately after downloading the model. It safely patches the code bugs in the HuggingFace cache and injects a mathematically safe PyTorch attention fallback so the model can generate motion perfectly without Flash Attention.

# Delete corrupted cache
import shutil
import glob
import sys
from transformers import AutoModel

shutil.rmtree("/root/.cache/huggingface/hub/models--ShandaAI--FloodDiffusionTiny", ignore_errors=True)
shutil.rmtree("/root/.cache/huggingface/modules/transformers_modules/ShandaAI", ignore_errors=True)
print("✅ Cache cleared")

# 1. Trigger first download (downloads hf_pipeline.py, then crashes)
try:
    AutoModel.from_pretrained("ShandaAI/FloodDiffusionTiny", trust_remote_code=True)
except Exception:
    pass

# 2. trust_remote_code patch
for f in glob.glob("/root/.cache/huggingface/modules/transformers_modules/ShandaAI/FloodDiffusionTiny/*/hf_pipeline.py"):
    text = open(f).read()
    text = text.replace(
        "if not kwargs.get('trust_remote_code', False):",
        "if False:"
    )
    open(f, 'w').write(text)
    print("✅ Patch 1 applied")

# 3. Trigger second download (downloads the rest of the model, then crashes on dtype)
try:
    AutoModel.from_pretrained("ShandaAI/FloodDiffusionTiny", trust_remote_code=True)
except Exception:
    pass

# 4. Patch the dtype bug
for f in glob.glob("/root/.cache/huggingface/hub/models--ShandaAI--FloodDiffusionTiny/*/*/ldf_models/diffusion_forcing_wan_tiny.py"):
    text = open(f).read()
    text = text.replace(
        "dtype=dtype\n        ).encoder",
        "torch_dtype=dtype\n        ).encoder"
    )
    open(f, 'w').write(text)
    print("✅ Patch 2 (dtype) applied")

# 5. Inject Safe PyTorch Attention fallback
# NOTE: You can skip this entire block (Step 5) if you successfully installed `flash-attn`
for f in glob.glob("/root/.cache/huggingface/hub/models--ShandaAI--FloodDiffusionTiny/*/*/ldf_models/tools/attention.py"):
    with open(f, 'r') as file:
        lines = file.readlines()

    else_idx = -1
    return_idx = -1

    # Find exactly where the else: block starts and where the function returns
    for i, line in enumerate(lines):
        if line.startswith("    else:"):
            else_idx = i
        if line.startswith("    # output") or line.startswith("    return"):
            if else_idx != -1:
                return_idx = i
                break

    if else_idx != -1 and return_idx != -1:
        new_lines = lines[:else_idx+1] # Keep everything up to 'else:'

        # A perfect SDPA fallback that correctly chunks the flattened tensors!
        new_code = """        import torch.nn.functional as _F
        out = []
        q_start = 0
        k_start = 0
        for i in range(b):
            q_len = q_lens[i].item()
            k_len = k_lens[i].item()
            q_i = q[q_start:q_start+q_len].transpose(0, 1) # [Nq, Lq, D]
            k_i = k[k_start:k_start+k_len].transpose(0, 1) # [Nk, Lk, D]
            v_i = v[k_start:k_start+k_len].transpose(0, 1)
            
            x_i = _F.scaled_dot_product_attention(
                q_i.unsqueeze(0), k_i.unsqueeze(0), v_i.unsqueeze(0),
                dropout_p=dropout_p if dropout_p else 0.0,
                is_causal=causal if causal else False,
                scale=softmax_scale,
            ).squeeze(0).transpose(0, 1) # [Lq, Nq, D]
            out.append(x_i)
            
            q_start += q_len
            k_start += k_len
        x = torch.cat(out, dim=0).unflatten(0, (b, lq))
"""
        new_lines.append(new_code)
        new_lines.extend(lines[return_idx:]) # Keep the return statement

        with open(f, 'w') as file:
            file.writelines(new_lines)
        print("✅ Patch 3 (Attention gracefully handles unpadded variable lengths!)")

# Clean memory
for key in list(sys.modules.keys()):
    if 'FloodDiffusion' in key or 'ldf' in key or 'hf_pipeline' in key or 'diffusion_forcing' in key or 'attention' in key:
        del sys.modules[key]


# === Run the Model ===
import torch

print("\nLoading model...")
model = AutoModel.from_pretrained(
    "ShandaAI/FloodDiffusionTiny",
    trust_remote_code=True
)
print("🎉 Model loaded successfully!")

model = model.to('cuda')

# Generate motion
print("🏃 Generating motion...")
motion = model("a person walking forward", length=60)
print(f"✅ Generated motion shape: {motion.shape}")