Update app.py

app.py

@@ -1,6 +1,6 @@
+import os
 import spaces
 import torch
-import os  # ← مضاف
 from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
 from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
 from diffusers.utils.export_utils import export_to_video
@@ -18,7 +18,8 @@ import aoti
 # =========================================================
 # MODEL CONFIGURATION
 # =========================================================
-MODEL_ID = "dream2589632147/Dream-wan2-2-faster-Pro"
+MODEL_ID = "dream2589632147/Dream-wan2-2-faster-Pro"  # المسار الجديد للنموذج
+HF_TOKEN = os.environ.get("HF_TOKEN")  # ضع توكن Hugging Face هنا إذا كان النموذج خاصًا
 
 MAX_DIM = 832
 MIN_DIM = 480
@@ -26,40 +27,34 @@ SQUARE_DIM = 640
 MULTIPLE_OF = 16
 
 MAX_SEED = np.iinfo(np.int32).max
+
 FIXED_FPS = 16
 MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 720
+MAX_FRAMES_MODEL = 80
+
 MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
 MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)
 
-# =========================================================
-# LOAD TRANSFORMERS (باستخدام HF_TOKEN)
-# =========================================================
-transformer = WanTransformer3DModel.from_pretrained(
-    MODEL_ID,
-    subfolder="transformer",
-    torch_dtype=torch.bfloat16,
-    device_map="cuda",
-    token=os.environ.get("HF_TOKEN")  # ← يستخدم التوكن للوصول إلى النموذج
-)
-
-transformer_2 = WanTransformer3DModel.from_pretrained(
-    MODEL_ID,
-    subfolder="transformer_2",
-    torch_dtype=torch.bfloat16,
-    device_map="cuda",
-    token=os.environ.get("HF_TOKEN")
-)
-
 # =========================================================
 # LOAD PIPELINE
 # =========================================================
 pipe = WanImageToVideoPipeline.from_pretrained(
     MODEL_ID,
-    transformer=transformer,
-    transformer_2=transformer_2,
+    transformer=WanTransformer3DModel.from_pretrained(
+        MODEL_ID,
+        subfolder="transformer",
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+        token=HF_TOKEN
+    ),
+    transformer_2=WanTransformer3DModel.from_pretrained(
+        MODEL_ID,
+        subfolder="transformer_2",
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+        token=HF_TOKEN
+    ),
     torch_dtype=torch.bfloat16,
-    token=os.environ.get("HF_TOKEN")
 ).to("cuda")
 
 # =========================================================
@@ -113,6 +108,7 @@ def resize_image(image: Image.Image) -> Image.Image:
     aspect_ratio = width / height
     MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
     MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM
+
     image_to_resize = image
 
     if aspect_ratio > MAX_ASPECT_RATIO:
@@ -133,12 +129,14 @@ def resize_image(image: Image.Image) -> Image.Image:
 
     final_w = round(target_w / MULTIPLE_OF) * MULTIPLE_OF
     final_h = round(target_h / MULTIPLE_OF) * MULTIPLE_OF
+
     final_w = max(MIN_DIM, min(MAX_DIM, final_w))
     final_h = max(MIN_DIM, min(MAX_DIM, final_h))
+
     return image_to_resize.resize((final_w, final_h), Image.LANCZOS)
 
 # =========================================================
-# UTILITIES
+# UTILITY FUNCTIONS
 # =========================================================
 def get_num_frames(duration_seconds: float):
     return 1 + int(np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL))
@@ -157,7 +155,7 @@ def get_duration(
     return 10 + int(steps) * step_duration
 
 # =========================================================
-# MAIN FUNCTION
+# MAIN GENERATION FUNCTION
 # =========================================================
 @spaces.GPU(duration=get_duration)
 def generate_video(