app.py

import os  # Для создания директории (на всякий случай оставляем)
import sys
import uuid
import shutil
import time
import gradio as gr
import torch
from diffusers import StableVideoDiffusionPipeline
from PIL import Image
import numpy as np
import cv2
import tempfile
from diffusers.utils import export_to_video  # Для экспорта видео

class WanAnimateApp:
    def __init__(self):
        model_name = "stabilityai/stable-video-diffusion-img2vid-xt"
        self.pipe = StableVideoDiffusionPipeline.from_pretrained(
            model_name,
            torch_dtype=torch.float32,  # Для CPU
            variant="fp16",
            low_cpu_mem_usage=True  # Оптимизация памяти для CPU
        )
        self.pipe.to("cpu")  # Вручную перемещаем на CPU

    def predict(self, ref_img, video, model_id, model):
        if ref_img is None or video is None:
            return None, "Upload both image and video."

        try:
            # Обработка изображения (теперь ref_img — np.array, конвертируем в PIL)
            ref_image = Image.fromarray(ref_img).convert("RGB").resize((576, 320))

            # Извлечение motion из видео (video — filepath)
            cap = cv2.VideoCapture(video)
            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            cap.release()
            motion_hint = f" with dynamic motion from {frame_count} frames"

            # Параметры
            num_frames = 25 if model == "wan-pro" else 14
            num_steps = 25 if model == "wan-pro" else 15

            # Адаптация modes
            noise_aug_strength = 0.02
            if model_id == "wan2.2-animate-mix":
                noise_aug_strength = 0.1

            # Генерация
            generator = torch.Generator(device="cpu").manual_seed(42)
            output = self.pipe(
                ref_image,
                num_inference_steps=num_steps,
                num_frames=num_frames,
                generator=generator,
                decode_chunk_size=2,  # Оптимизация для VAE
                noise_aug_strength=noise_aug_strength
            ).frames[0]

            # Экспорт видео
            with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
                export_to_video(output, temp_video.name, fps=7)

            return temp_video.name, "SUCCEEDED" + motion_hint

        except Exception as e:
            return None, f"Failed: {str(e)}"

def start_app():
    # Создаём директорию для Gradio (на всякий случай)
    os.makedirs("/tmp/gradio", exist_ok=True)
    
    app = WanAnimateApp()
    with gr.Blocks(title="Wan2.2-Animate (Local No API)") as demo:
        gr.HTML("""
            Wan2.2-Animate: Unified Character Animation and Replacement with Holistic Replication
            Local version without API (SVD Proxy)
            Tongyi Lab, Alibaba
            📄Paper 💻GitHub 🤗HF Model
        """)

        gr.HTML("""
            ‼️Usage (использования) Wan-Animate supports two modes:

            * Move Mode: animate the character in input image with movements from the input video

            * Mix Mode: replace the character in input video with the character in input image

            Wan-Animate supports two modes:

            * Move Mode: Use the movements extracted from the input video to drive the character in the input image

            * Mix Mode: Use the character in the input image to replace the character in the input video

            Currently, the following restrictions apply to inputs:

            * Video file size: Less than 200MB

            * Video resolution: The shorter side must be greater than 200, and the longer side must be less than 2048

            * Video duration: 2s to 30s

            * Video aspect ratio: 1:3 to 3:1

            * Video formats: mp4, avi, mov

            * Image file size: Less than 5MB

            * Image resolution: The shorter side must be greater than 200, and the longer side must be less than 4096

            * Image formats: jpg, png, jpeg, webp, bmp

            Current, the inference quality has two variants. You can use our open-source code for more flexible configuration.

            * wan-pro: 25fps, 720p

            * wan-std: 15fps, 720p
        """)

        with gr.Row():
            with gr.Column():
                ref_img = gr.Image(label="Reference Image (изображение)", type="numpy", sources=["upload"])  # Изменили на numpy для обхода FileNotFound
                video = gr.Video(label="Template Video (шаблонное видео)", sources=["upload"])
                with gr.Row():
                    model_id = gr.Dropdown(label="Mode (режим)", choices=["wan2.2-animate-move", "wan2.2-animate-mix"], value="wan2.2-animate-move")
                    model = gr.Dropdown(label="Inference Quality (качество)", choices=["wan-pro", "wan-std"], value="wan-pro")
                run_button = gr.Button("Generate Video (генерировать)")

            with gr.Column():
                output_video = gr.Video(label="Output Video (результат)")
                output_status = gr.Textbox(label="Status (статус)")

        run_button.click(
            fn=app.predict,
            inputs=[ref_img, video, model_id, model],
            outputs=[output_video, output_status]
        )

    demo.queue(default_concurrency_limit=1)
    demo.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    start_app()