Update app.py

app.py

@@ -35,9 +35,9 @@ MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)
 # =========================================================
 # LOAD PIPELINE
 # =========================================================
-print("Loading pipeline...")
+print("Loading pipeline components...")
 
-# 1. تحميل المكونات بدون تحديد device_map="cuda" لتجنب تعارض ZeroGPU
+# تحميل المكونات أولاً بدون نقلها للـ GPU لتوفير الذاكرة أثناء التحميل
 transformer = WanTransformer3DModel.from_pretrained(
     MODEL_ID,
     subfolder="transformer",
@@ -52,51 +52,62 @@ transformer_2 = WanTransformer3DModel.from_pretrained(
     token=HF_TOKEN
 )
 
-# 2. تجميع البايبلاين
+print("Assembling pipeline...")
 pipe = WanImageToVideoPipeline.from_pretrained(
     MODEL_ID,
     transformer=transformer,
     transformer_2=transformer_2,
     torch_dtype=torch.bfloat16,
+    token=HF_TOKEN
 )
 
-# 3. نقل الموديل للـ CUDA مرة واحدة هنا
-print("Moving pipeline to CUDA...")
+# نقل الموديل إلى CUDA الآن
+print("Moving to CUDA...")
 pipe = pipe.to("cuda")
 
 # =========================================================
 # LOAD LORA ADAPTERS
 # =========================================================
 print("Loading LoRA adapters...")
-pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy",
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-    adapter_name="lightx2v"
-)
-pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy",
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
-    adapter_name="lightx2v_2",
-    load_into_transformer_2=True
-)
+try:
+    pipe.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+        adapter_name="lightx2v"
+    )
+    pipe.load_lora_weights(
+        "Kijai/WanVideo_comfy",
+        weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
+        adapter_name="lightx2v_2",
+        load_into_transformer_2=True
+    )
 
-pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
-pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
-pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
-pipe.unload_lora_weights()
+    pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
+    pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
+    pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
+    pipe.unload_lora_weights()
+    print("LoRA loaded and fused successfully.")
+except Exception as e:
+    print(f"Warning: Failed to load LoRA. Continuing without it. Error: {e}")
 
 # =========================================================
 # QUANTIZATION & AOT OPTIMIZATION
 # =========================================================
 print("Applying quantization...")
-# نقلنا التكميم بعد النقل للـ GPU لضمان التوافق
-quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
-quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
-quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+# تنظيف الذاكرة قبل العمليات الثقيلة
+torch.cuda.empty_cache()
+gc.collect()
 
-print("Loading AOTI blocks...")
-aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
-aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
+try:
+    quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+    quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+    quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+
+    print("Loading AOTI blocks...")
+    aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
+    aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
+except Exception as e:
+    print(f"Warning: Quantization/AOTI failed. Running in standard mode might OOM. Error: {e}")
 
 # =========================================================
 # DEFAULT PROMPTS
@@ -150,26 +161,11 @@ def resize_image(image: Image.Image) -> Image.Image:
 def get_num_frames(duration_seconds: float):
     return 1 + int(np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL))
 
-def get_duration(
-    input_image, prompt, steps, negative_prompt,
-    duration_seconds, guidance_scale, guidance_scale_2,
-    seed, randomize_seed, progress,
-):
-    if input_image is None:
-        return 120
-    BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
-    BASE_STEP_DURATION = 15
-    
-    width, height = resize_image(input_image).size
-    frames = get_num_frames(duration_seconds)
-    factor = frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
-    step_duration = BASE_STEP_DURATION * factor ** 1.5
-    return 10 + int(steps) * step_duration
-
 # =========================================================
 # MAIN GENERATION FUNCTION
 # =========================================================
-@spaces.GPU(duration=get_duration)
+# زيادة الوقت المسموح به إلى 180 ثانية لتجنب التايم أوت
+@spaces.GPU(duration=180)
 def generate_video(
     input_image,
     prompt,
@@ -182,31 +178,49 @@ def generate_video(
     randomize_seed=False,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if input_image is None:
-        raise gr.Error("Please upload an input image.")
-        
-    num_frames = get_num_frames(duration_seconds)
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    resized_image = resize_image(input_image)
+    # تنظيف الذاكرة في بداية الدالة
+    gc.collect()
+    torch.cuda.empty_cache()
+    
+    try:
+        if input_image is None:
+            raise gr.Error("Please upload an input image.")
+            
+        num_frames = get_num_frames(duration_seconds)
+        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+        resized_image = resize_image(input_image)
 
-    output_frames_list = pipe(
-        image=resized_image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=resized_image.height,
-        width=resized_image.width,
-        num_frames=num_frames,
-        guidance_scale=float(guidance_scale),
-        guidance_scale_2=float(guidance_scale_2),
-        num_inference_steps=int(steps),
-        generator=torch.Generator(device="cuda").manual_seed(current_seed),
-    ).frames[0]
+        print(f"Generating video with seed: {current_seed}, frames: {num_frames}")
 
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
-        video_path = tmpfile.name
-    
-    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
-    return video_path, current_seed
+        output_frames_list = pipe(
+            image=resized_image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=resized_image.height,
+            width=resized_image.width,
+            num_frames=num_frames,
+            guidance_scale=float(guidance_scale),
+            guidance_scale_2=float(guidance_scale_2),
+            num_inference_steps=int(steps),
+            generator=torch.Generator(device="cuda").manual_seed(current_seed),
+        ).frames[0]
+
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+            video_path = tmpfile.name
+        
+        export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
+        
+        # تنظيف الذاكرة بعد الانتهاء
+        del output_frames_list
+        torch.cuda.empty_cache()
+        
+        return video_path, current_seed
+
+    except Exception as e:
+        # طباعة الخطأ الحقيقي في الكونسول
+        print(f"Error during generation: {e}")
+        # إعادة رفع الخطأ ليظهر للمستخدم
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 # =========================================================
 # GRADIO UI
@@ -285,9 +299,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
-    # ملاحظة: يمكنك حذف جزء Examples إذا استمرت الأخطاء أو إذا لم تكن الصورة مرفوعة
-    # gr.Examples(...)
-
     # --- BOTTOM ADVERTISEMENT BANNER ---
     gr.HTML("""
     <div style="background: linear-gradient(90deg, #4f46e5, #9333ea); color: white; padding: 15px; border-radius: 10px; text-align: center; margin-top: 20px; box-shadow: 0 4px 15px rgba(0,0,0,0.1);">
@@ -306,4 +317,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     """)
 
 if __name__ == "__main__":
-    demo.queue().launch(mcp_server=True)
+    # تم إزالة mcp_server=True لأنه يسبب مشاكل
+    demo.queue().launch()