update

app.py

@@ -6,6 +6,11 @@ import subprocess
 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"])
+
+# Cài đặt ffmpeg trong môi trường Ubuntu
+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'):
@@ -16,6 +21,78 @@ if not os.path.exists('first_order_model'):
 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("""
+import yaml
+import torch
+from first_order_model.modules.generator import OcclusionAwareGenerator
+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)
+        
+    generator = OcclusionAwareGenerator(**config['model_params']['generator_params'],
+                                        **config['model_params']['common_params'])
+    generator.to(device)
+    
+    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
+                              **config['model_params']['common_params'])
+    kp_detector.to(device)
+    
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    generator.load_state_dict(checkpoint['generator'])
+    kp_detector.load_state_dict(checkpoint['kp_detector'])
+    
+    generator.eval()
+    kp_detector.eval()
+    
+    return generator, kp_detector
+    
+def normalize_kp(kp_source, kp_driving, kp_driving_initial, 
+                 use_relative_movement=True, use_relative_jacobian=True, adapt_movement_scale=True):
+    from first_order_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+    
+    kp_new = {k: v for k, v in kp_driving.items()}
+    
+    if use_relative_movement:
+        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
+        kp_value_diff_abs = torch.abs(kp_value_diff)
+        
+        if adapt_movement_scale:
+            distance = torch.max(kp_value_diff_abs, dim=2, keepdim=True)[0]
+            distance = torch.max(distance, dim=1, keepdim=True)[0]
+            
+            kp_source_diff = torch.abs(kp_source['value'])
+            kp_source_max = torch.max(kp_source_diff, dim=2, keepdim=True)[0]
+            kp_source_max = torch.max(kp_source_max, dim=1, keepdim=True)[0]
+            
+            movement_scale = kp_source_max / (distance + 1e-6)
+            
+            kp_new['value'] = kp_source['value'] + movement_scale * kp_value_diff
+        else:
+            kp_new['value'] = kp_source['value'] + kp_value_diff
+            
+    if use_relative_jacobian:
+        jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))
+        kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])
+        
+    return kp_new
+""")
+
 # Bây giờ import các module cần thiết
 import gradio as gr
 import numpy as np
@@ -23,28 +100,26 @@ import torch
 import imageio
 from skimage.transform import resize
 from skimage import img_as_ubyte
+from PIL import Image
 
-# Import các module từ first_order_model
-from first_order_model.demo import load_checkpoints
-from first_order_model.animate import normalize_kp
+# Import hàm load_checkpoints từ file helper
+from load_helper import load_checkpoints, normalize_kp
 
 # Tải mô hình pre-trained
 def download_model():
     model_path = 'checkpoints/vox-cpk.pth.tar'
-    if not os.path.exists(model_path):
+    if not os.path.exists('checkpoints'):
         os.makedirs('checkpoints', exist_ok=True)
-        subprocess.call([
-            'wget', 'https://drive.google.com/uc?export=download&id=1PyQJmkdCsAkOYwUyaj_l-l0as-iLDgeH', 
-            '-O', model_path
-        ])
+    
+    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')
     
     config_path = 'first_order_model/config/vox-256.yaml'
-    if not os.path.exists(config_path):
+    if not os.path.exists('first_order_model/config'):
         os.makedirs('first_order_model/config', exist_ok=True)
-        subprocess.call([
-            'wget', 'https://drive.google.com/uc?export=download&id=1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn',
-            '-O', config_path
-        ])
+    
+    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')
     
     return config_path, model_path
 
@@ -53,6 +128,7 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
     config_path, checkpoint_path = download_model()
     
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
     
     # Tải mô hình và cấu hình
     generator, kp_detector = load_checkpoints(config_path, checkpoint_path, device=device)
@@ -89,17 +165,14 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
             kp_driving = kp_detector(driving_frame)
             
             # Chuẩn hóa keypoints
-            if relative:
-                kp_norm = normalize_kp(
-                    kp_source=kp_source,
-                    kp_driving=kp_driving,
-                    kp_driving_initial=kp_driving_initial,
-                    use_relative_movement=relative,
-                    use_relative_jacobian=relative,
-                    adapt_movement_scale=adapt_movement_scale
-                )
-            else:
-                kp_norm = kp_driving
+            kp_norm = normalize_kp(
+                kp_source=kp_source,
+                kp_driving=kp_driving,
+                kp_driving_initial=kp_driving_initial,
+                use_relative_movement=relative,
+                use_relative_jacobian=relative,
+                adapt_movement_scale=adapt_movement_scale
+            )
             
             # Tạo frame
             out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
@@ -107,7 +180,11 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
     
     # Lưu video kết quả
     output_path = 'result.mp4'
-    imageio.mimsave(output_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)
+    os.system(f"rm -f {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]
+    imageio.mimsave(output_path, frames, fps=fps)
     
     return output_path
 

requirements.txt

@@ -8,4 +8,5 @@ scikit-image==0.19.3
 scikit-learn
 matplotlib
 PyYAML==5.3.1
-face-alignment==1.3.5
+face-alignment==1.3.5
+ffmpeg-python