update

app.py

@@ -1,13 +1,21 @@
 import os
 import sys
 import subprocess
+import numpy as np
+import torch
+import imageio
+from skimage.transform import resize
+from skimage import img_as_ubyte
+import gradio as gr
+from PIL import Image
+import tempfile
 
 # Đảm bảo cài đặt các thư viện cần thiết trước khi import
 subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-learn"])
 subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-image==0.19.3"])
 subprocess.check_call([sys.executable, "-m", "pip", "install", "face-alignment==1.3.5"])
-subprocess.check_call([sys.executable, "-m", "pip", "install", "ffmpeg-python"])
 subprocess.check_call([sys.executable, "-m", "pip", "install", "PyYAML==5.3.1"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "imageio-ffmpeg==0.4.5"])
 
 # Cài đặt ffmpeg trong môi trường Ubuntu
 os.system("apt-get update && apt-get install -y ffmpeg")
@@ -15,24 +23,13 @@ os.system("apt-get update && apt-get install -y ffmpeg")
 # Clone repo nếu chưa có
 if not os.path.exists('first_order_model'):
     subprocess.call(['git', 'clone', 'https://github.com/AliaksandrSiarohin/first-order-model.git'])
-    os.rename('first-order-model', 'first_order_model')
+    if os.path.exists('first-order-model'):
+        os.rename('first-order-model', 'first_order_model')
 
 # Thêm đường dẫn vào PYTHONPATH
 sys.path.append('.')
 sys.path.append('first_order_model')
 
-# Sửa code để truy cập trực tiếp vào các hàm cần thiết
-# Tạo một bản sao của demo.py mà không phụ thuộc vào ffmpeg thư viện Python
-with open('first_order_model/demo.py', 'r') as f:
-    demo_code = f.read()
-
-# Thay thế dòng import ffmpeg
-demo_code = demo_code.replace('import ffmpeg', '# import ffmpeg')
-
-# Viết lại demo.py đã sửa
-with open('first_order_model/demo_fixed.py', 'w') as f:
-    f.write(demo_code)
-
 # Tạo file helper với hàm load_checkpoints
 with open('load_helper.py', 'w') as f:
     f.write("""
@@ -43,7 +40,7 @@ from first_order_model.modules.keypoint_detector import KPDetector
 
 def load_checkpoints(config_path, checkpoint_path, device='cpu'):
     with open(config_path) as f:
-        config = yaml.full_load(f)
+        config = yaml.safe_load(f)
         
     generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
                                         **config['model_params']['common_params'])
@@ -93,15 +90,6 @@ def normalize_kp(kp_source, kp_driving, kp_driving_initial,
     return kp_new
 """)
 
-# Bây giờ import các module cần thiết
-import gradio as gr
-import numpy as np
-import torch
-import imageio
-from skimage.transform import resize
-from skimage import img_as_ubyte
-from PIL import Image
-
 # Import hàm load_checkpoints từ file helper
 from load_helper import load_checkpoints, normalize_kp
 
@@ -112,14 +100,14 @@ def download_model():
         os.makedirs('checkpoints', exist_ok=True)
     
     if not os.path.exists(model_path):
-        os.system('wget -P checkpoints https://drive.google.com/uc?export=download&id=1PyQJmkdCsAkOYwUyaj_l-l0as-iLDgeH -O checkpoints/vox-cpk.pth.tar')
+        os.system('wget -O checkpoints/vox-cpk.pth.tar https://drive.google.com/uc?export=download&id=1PyQJmkdCsAkOYwUyaj_l-l0as-iLDgeH')
     
     config_path = 'first_order_model/config/vox-256.yaml'
     if not os.path.exists('first_order_model/config'):
         os.makedirs('first_order_model/config', exist_ok=True)
     
     if not os.path.exists(config_path):
-        os.system('wget -P first_order_model/config https://drive.google.com/uc?export=download&id=1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn -O first_order_model/config/vox-256.yaml')
+        os.system('wget -O first_order_model/config/vox-256.yaml https://drive.google.com/uc?export=download&id=1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn')
     
     return config_path, model_path
 
@@ -180,7 +168,8 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
     
     # Lưu video kết quả
     output_path = 'result.mp4'
-    os.system(f"rm -f {output_path}")  # Xóa video nếu tồn tại
+    if os.path.exists(output_path):
+        os.remove(output_path)  # Xóa video nếu tồn tại
     
     # Lưu frames thành video sử dụng imageio
     frames = [img_as_ubyte(frame) for frame in predictions]
@@ -189,21 +178,34 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
     return output_path
 
 # Định nghĩa giao diện Gradio
-def animate_fomm(source_image, driving_video, relative=True, adapt_scale=True):
-    if source_image is None or driving_video is None:
+def animate_fomm(source_image, driving_video_file, relative=True, adapt_scale=True):
+    if source_image is None or driving_video_file is None:
         return None, "Vui lòng tải lên cả ảnh nguồn và video tham chiếu."
     
     try:
-        # Lưu tạm ảnh và video tải lên
+        # Lưu tạm ảnh nguồn
         source_path = "source_image.jpg"
-        driving_path = "driving_video.mp4"
-        
-        # Lưu ảnh nguồn
         source_image.save(source_path)
         
-        # Lưu video tham chiếu
-        with open(driving_path, 'wb') as f:
-            f.write(driving_video)
+        # Xử lý video tham chiếu - Sửa lỗi "a bytes-like object is required, not 'str'"
+        print(f"Type of driving_video: {type(driving_video_file)}")
+        
+        # Tạo file tạm cho video
+        driving_path = "driving_video.mp4"
+        
+        # Kiểm tra loại dữ liệu của driving_video_file
+        if isinstance(driving_video_file, str):
+            # Nếu là đường dẫn, sử dụng video mẫu hoặc copy file
+            if os.path.exists(driving_video_file):
+                import shutil
+                shutil.copyfile(driving_video_file, driving_path)
+            else:
+                # Tải video mẫu
+                os.system("wget -O driving_video.mp4 https://github.com/AliaksandrSiarohin/first-order-model/raw/master/driving.mp4")
+        else:
+            # Ghi dữ liệu nhị phân vào file
+            with open(driving_path, 'wb') as f:
+                f.write(driving_video_file)
         
         # Tạo animation
         result_path = make_animation(
@@ -215,7 +217,8 @@ def animate_fomm(source_image, driving_video, relative=True, adapt_scale=True):
         
         return result_path, "Video được tạo thành công!"
     except Exception as e:
-        return None, f"Lỗi: {str(e)}"
+        import traceback
+        return None, f"Lỗi: {str(e)}\n{traceback.format_exc()}"
 
 # Tạo giao diện Gradio
 with gr.Blocks(title="First Order Motion Model - Tạo video người chuyển động") as demo:
@@ -225,28 +228,51 @@ with gr.Blocks(title="First Order Motion Model - Tạo video người chuyển
     with gr.Row():
         with gr.Column():
             source_image = gr.Image(type="pil", label="Tải lên ảnh nguồn")
-            driving_video = gr.Video(label="Tải lên video tham chiếu")
+            
+            # Thay đổi từ gr.Video sang gr.File để xử lý lỗi binary
+            driving_video_file = gr.File(label="Tải lên video tham chiếu (.mp4)")
             
             with gr.Row():
                 relative = gr.Checkbox(value=True, label="Chuyển động tương đối")
                 adapt_scale = gr.Checkbox(value=True, label="Điều chỉnh tỷ lệ chuyển động")
             
+            # Thêm tùy chọn sử dụng video mẫu
+            use_sample = gr.Checkbox(label="Sử dụng video mẫu (nếu bạn không muốn tải lên video)")
+            
             submit_btn = gr.Button("Tạo video")
             
         with gr.Column():
             output_video = gr.Video(label="Video kết quả")
-            output_message = gr.Textbox(label="Thông báo")
+            output_message = gr.Textbox(label="Thông báo", lines=5)
+    
+    # Xử lý sự kiện khi checkbox được chọn
+    def toggle_video_upload(use_sample_video):
+        return gr.update(interactive=not use_sample_video)
+    
+    use_sample.change(fn=toggle_video_upload, inputs=[use_sample], outputs=[driving_video_file])
+    
+    # Cập nhật hàm xử lý khi nhấn nút
+    def process_inputs(source_img, driving_vid, use_sample_vid, rel, adapt):
+        if use_sample_vid:
+            # Tải video mẫu nếu cần
+            sample_path = "sample_driving.mp4"
+            if not os.path.exists(sample_path):
+                os.system("wget -O sample_driving.mp4 https://github.com/AliaksandrSiarohin/first-order-model/raw/master/driving.mp4")
+            return animate_fomm(source_img, sample_path, rel, adapt)
+        else:
+            return animate_fomm(source_img, driving_vid, rel, adapt)
     
     submit_btn.click(
-        fn=animate_fomm,
-        inputs=[source_image, driving_video, relative, adapt_scale],
+        fn=process_inputs,
+        inputs=[source_image, driving_video_file, use_sample, relative, adapt_scale],
         outputs=[output_video, output_message]
     )
     
     gr.Markdown("### Cách sử dụng")
     gr.Markdown("1. Tải lên **ảnh nguồn** - ảnh chứa người/đối tượng bạn muốn làm chuyển động")
     gr.Markdown("2. Tải lên **video tham chiếu** - video có chuyển động bạn muốn áp dụng")
-    gr.Markdown("3. Nhấn **Tạo video** và chờ kết quả")
+    gr.Markdown("3. Hoặc chọn sử dụng video mẫu có sẵn")
+    gr.Markdown("4. Nhấn **Tạo video** và chờ kết quả")
     
     gr.Markdown("### Lưu ý")
     gr.Markdown("- Ảnh nguồn và video tham chiếu nên có đối tượng tương tự (người với người, mặt với mặt)")