Update app.py

app.py

@@ -12,6 +12,14 @@ import cv2
 import subprocess
 import tempfile
 
+# Папки (создаются автоматически)
+IMG_DIR = './img'
+VIDEO_DIR = './video'
+OUTPUT_DIR = './output'
+os.makedirs(IMG_DIR, exist_ok=True)
+os.makedirs(VIDEO_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
 class WanAnimateApp:
     def __init__(self):
         model_name = "stabilityai/stable-video-diffusion-img2vid-xt"
@@ -22,25 +30,49 @@ class WanAnimateApp:
             device_map="cpu"
         )
 
+    def get_img_files(self):
+        """Список файлов в папке img"""
+        if os.path.exists(IMG_DIR):
+            return [f for f in os.listdir(IMG_DIR) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        return []
+
+    def get_video_files(self):
+        """Список файлов в папке video"""
+        if os.path.exists(VIDEO_DIR):
+            return [f for f in os.listdir(VIDEO_DIR) if f.lower().endswith(('.mp4', '.avi', '.mov'))]
+        return []
+
+    def load_img_from_dropdown(self, selected_file):
+        """Загружает выбранное фото в ref_img"""
+        if selected_file:
+            path = os.path.join(IMG_DIR, selected_file)
+            if os.path.exists(path):
+                return path
+        return None
+
+    def load_video_from_dropdown(self, selected_file):
+        """Загружает выбранное видео в video"""
+        if selected_file:
+            path = os.path.join(VIDEO_DIR, selected_file)
+            if os.path.exists(path):
+                return path
+        return None
+
     def predict(
         self, 
-        ref_img,
-        video,
+        ref_img_path,
+        video_path,
         model_id,
         model,
     ):
-        if ref_img is None or video is None:
-            return None, "Upload both image and video."
+        if ref_img_path is None or video_path is None:
+            return None, "Выберите файлы из dropdown."
 
         try:
-            # Local processing — PIL for image (no open)
-            if isinstance(ref_img, Image.Image):
-                ref_image = ref_img.convert("RGB").resize((576, 320))
-            else:
-                ref_image = Image.open(ref_img).convert("RGB").resize((576, 320))
-
-            cap = cv2.VideoCapture(video)
-            frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            # Local processing (пути — реальные str из папок)
+            ref_image = Image.open(ref_img_path).convert("RGB").resize((576, 320))
+            cap = cv2.VideoCapture(video_path)
+            frame_count = Image.open(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             cap.release()
             motion_hint = f" with dynamic motion from {frame_count} frames"
 
@@ -68,7 +100,7 @@ class WanAnimateApp:
             temp_dir = tempfile.mkdtemp()
             for i, frame in enumerate(output):
                 frame.save(f"{temp_dir}/frame_{i:04d}.png")
-            temp_video = f"/tmp/output_{uuid.uuid4()}.mp4"
+            temp_video = os.path.join(OUTPUT_DIR, f"output_{uuid.uuid4()}.mp4")
             subprocess.run([
                 'ffmpeg', '-y', '-framerate', '7', '-i', f"{temp_dir}/frame_%04d.png",
                 '-c:v', 'libx264', '-pix_fmt', 'yuv420p', temp_video
@@ -143,15 +175,38 @@ def start_app():
 
         with gr.Row():
             with gr.Column():    
+                # Dropdown for img
+                img_dropdown = gr.Dropdown(
+                    label="Select Image from img folder (выбери фото)",
+                    choices=app.get_img_files(),
+                    value=None,
+                )
                 ref_img = gr.Image(
                     label="Reference Image (изображение)",
-                    type="pil",
-                    sources=["upload"],
+                    type="filepath",
+                    value=None,
+                )
+                img_dropdown.change(
+                    fn=app.load_img_from_dropdown,
+                    inputs=[img_dropdown],
+                    outputs=[ref_img],
+                )
+
+                # Dropdown for video
+                video_dropdown = gr.Dropdown(
+                    label="Select Video from video folder (выбери видео)",
+                    choices=app.get_video_files(),
+                    value=None,
                 )
-                
                 video = gr.Video(
                     label="Template Video (шаблонное видео)",
-                    sources=["upload"],
+                    type="filepath",
+                    value=None,
+                )
+                video_dropdown.change(
+                    fn=app.load_video_from_dropdown,
+                    inputs=[video_dropdown],
+                    outputs=[video],
                 )
                 
                 with gr.Row():