import os

# Patch wan/__init__.py
p = '/app/Wan2.2/wan/__init__.py'
with open(p) as f:
    content = f.read()
lines = content.split('\n')
result = []
problematic = {'speech2video', 'animate', 'video2audio'}
for line in lines:
    stripped = line.strip()
    if stripped.startswith('from .') and 'import' in stripped:
        parts = stripped.split()
        module_ref = parts[1]
        module_name = module_ref.lstrip('.')
        if any(mod in module_name for mod in problematic):
            indent = line[:len(line) - len(line.lstrip())]
            result.append(f'{indent}try:')
            result.append(f'{indent}    {stripped}')
            result.append(f'{indent}except ImportError:')
            result.append(f'{indent}    pass')
        else:
            result.append(line)
    else:
        result.append(line)
with open(p, 'w') as f:
    f.write('\n'.join(result))
print('Patched wan/__init__.py')

# Patch generate.py
p = '/app/Wan2.2/generate.py'
with open(p) as f:
    c = f.read()
c = c.replace(
    'from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander',
    'try:\n    from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander\nexcept ImportError:\n    DashScopePromptExpander = None\n    QwenPromptExpander = None'
)
with open(p, 'w') as f:
    f.write(c)
print('Patched generate.py')