Update app.py

app.py

@@ -32,7 +32,16 @@ from PIL import Image
 from moviepy.editor import VideoFileClip, AudioFileClip
 from pydub import AudioSegment
 from huggingface_hub import snapshot_download
-from torchao.quantization import quantize_, int8_weight_only
+
+# ===== تغییرات اعمال شده برای رفع خطای torchao =====
+from torchao.quantization import quantize_
+try:
+    from torchao.quantization import Int8WeightOnlyConfig
+    quant_config = Int8WeightOnlyConfig()
+except ImportError:
+    from torchao.quantization import int8_weight_only
+    quant_config = int8_weight_only()
+# =================================================
 
 from src.models.unet_2d_condition import UNet2DConditionModel
 from src.models.unet_3d_emo import EMOUNet3DConditionModel
@@ -364,7 +373,8 @@ def build_pipeline(quantized=False, target_device=None, target_dtype=None):
         target_device, dtype=target_dtype
     )
     if quantized:
-        quantize_(vae, int8_weight_only())
+        # ===== تغییر مربوط به VAE =====
+        quantize_(vae, quant_config)
         print("Using int8 quantization for VAE.")
 
     reference_unet = UNet2DConditionModel.from_pretrained(
@@ -380,7 +390,8 @@ def build_pipeline(quantized=False, target_device=None, target_dtype=None):
         )
     )
     if quantized:
-        quantize_(reference_unet, int8_weight_only())
+        # ===== تغییر مربوط به reference UNet =====
+        quantize_(reference_unet, quant_config)
         print("Using int8 quantization for reference UNet.")
 
     motion_module_path = "./pretrained_weights/motion_module.pth"
@@ -1088,4 +1099,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.queue()
-    demo.launch(show_error=True, ssr_mode=False, css=css)
+    demo.launch(show_error=True, ssr_mode=False, css=css)