Update app.py

app.py

@@ -13,6 +13,26 @@ import uuid
 import shutil
 
 import importlib, site, sys
+from huggingface_hub import hf_hub_download, snapshot_download
+
+# Re-discover all .pth/.egg-link files
+for sitedir in site.getsitepackages():
+    site.addsitedir(sitedir)
+
+# Clear caches so importlib will pick up new modules
+importlib.invalidate_caches()
+
+def sh(cmd): subprocess.check_call(cmd, shell=True)
+
+flash_attention_wheel = hf_hub_download(
+        repo_id="alexnasa/flash-attn-3",
+        filename="flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl",
+    )
+
+sh(f"pip install {flash_attention_wheel}")
+
+# tell Python to re-scan site-packages now that the egg-link exists
+import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
 
 import torch
 
@@ -45,7 +65,7 @@ from transformers import Wav2Vec2FeatureExtractor
 import torchvision.transforms as transforms
 import torch.nn.functional as F
 from OmniAvatar.utils.audio_preprocess import add_silence_to_audio_ffmpeg
-from huggingface_hub import hf_hub_download, snapshot_download
+
 
 os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/proprocess_results"