update

app.py

@@ -16,6 +16,8 @@ subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-image==0.
 subprocess.check_call([sys.executable, "-m", "pip", "install", "face-alignment==1.3.5"])
 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"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "gdown"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
 
 # Cài đặt ffmpeg trong môi trường Ubuntu
 os.system("apt-get update && apt-get install -y ffmpeg")
@@ -93,23 +95,77 @@ def normalize_kp(kp_source, kp_driving, kp_driving_initial,
 # Import hàm load_checkpoints từ file helper
 from load_helper import load_checkpoints, normalize_kp
 
-# Tải mô hình pre-trained
+# Tải mô hình pre-trained với phương pháp cải tiến
 def download_model():
-    model_path = 'checkpoints/vox-cpk.pth.tar'
-    if not os.path.exists('checkpoints'):
-        os.makedirs('checkpoints', exist_ok=True)
-    
-    if not os.path.exists(model_path):
-        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 -O first_order_model/config/vox-256.yaml https://drive.google.com/uc?export=download&id=1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn')
+    try:
+        # Thử phương pháp sử dụng gdown trước
+        model_path = 'checkpoints/vox-cpk.pth.tar'
+        if not os.path.exists('checkpoints'):
+            os.makedirs('checkpoints', exist_ok=True)
+        
+        # Kiểm tra xem file đã tồn tại và đủ lớn chưa
+        if not os.path.exists(model_path) or os.path.getsize(model_path) < 1000000:
+            print("Đang tải mô hình từ Google Drive...")
+            import gdown
+            file_id = '1PyQJmkdCsAkOYwUyaj_l-l0as-iLDgeH'
+            gdown.download(f"https://drive.google.com/uc?id={file_id}", model_path, quiet=False)
+        
+        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) or os.path.getsize(config_path) < 1000:
+            print("Đang tải file cấu hình từ Google Drive...")
+            import gdown
+            file_id = '1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn'
+            gdown.download(f"https://drive.google.com/uc?id={file_id}", config_path, quiet=False)
+        
+        # Kiểm tra lại kích thước file
+        if os.path.exists(model_path) and os.path.getsize(model_path) > 1000000 and os.path.exists(config_path) and os.path.getsize(config_path) > 1000:
+            return config_path, model_path
+        
+        # Nếu gdown không thành công, thử phương pháp sử dụng Hugging Face
+        raise Exception("Tải mô hình từ Google Drive không thành công. Chuyển sang sử dụng Hugging Face...")
     
-    return config_path, model_path
+    except Exception as e:
+        print(f"Lỗi khi tải từ Google Drive: {str(e)}")
+        return download_model_alternative()
+
+def download_model_alternative():
+    try:
+        print("Đang tải mô hình từ Hugging Face...")
+        from huggingface_hub import hf_hub_download
+        
+        model_path = hf_hub_download(
+            repo_id="ycyunwei/first-order-motion-model", 
+            filename="vox-cpk.pth.tar"
+        )
+        
+        config_path = hf_hub_download(
+            repo_id="ycyunwei/first-order-motion-model",
+            filename="vox-256.yaml"
+        )
+        
+        return config_path, model_path
+    except Exception as e:
+        print(f"Lỗi khi tải từ Hugging Face: {str(e)}")
+        # Thử phương pháp cuối cùng - sử dụng direct link
+        model_path = 'checkpoints/vox-cpk.pth.tar'
+        if not os.path.exists('checkpoints'):
+            os.makedirs('checkpoints', exist_ok=True)
+        
+        os.system(f'wget -O {model_path} https://github.com/AliaksandrSiarohin/first-order-model/releases/download/v1.0.0/vox-cpk.pth.tar')
+        
+        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)
+        
+        os.system(f'wget -O {config_path} https://raw.githubusercontent.com/AliaksandrSiarohin/first-order-model/master/config/vox-256.yaml')
+        
+        if os.path.exists(model_path) and os.path.getsize(model_path) > 1000000 and os.path.exists(config_path) and os.path.getsize(config_path) > 1000:
+            return config_path, model_path
+        else:
+            raise Exception("Không thể tải mô hình b���ng cả ba phương pháp. Vui lòng thử lại sau.")
 
 # Hàm tạo animation
 def make_animation(source_image, driving_video, relative=True, adapt_movement_scale=True):
@@ -167,7 +223,7 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
             predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
     
     # Lưu video kết quả
-    output_path = 'result.mp4'
+    output_path = f'result_{int(np.random.rand() * 10000)}.mp4'
     if os.path.exists(output_path):
         os.remove(output_path)  # Xóa video nếu tồn tại
     
@@ -179,33 +235,41 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
 
 # Định nghĩa giao diện Gradio
 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."
+    if source_image is None:
+        return None, "Vui lòng tải lên ảnh nguồn."
     
     try:
         # Lưu tạm ảnh nguồn
-        source_path = "source_image.jpg"
+        source_path = f"source_image_{int(np.random.rand() * 10000)}.jpg"
         source_image.save(source_path)
         
-        # Xử lý video tham chiếu - Sửa lỗi "a bytes-like object is required, not 'str'"
+        # Xử lý video tham chiếu
         print(f"Type of driving_video: {type(driving_video_file)}")
         
         # Tạo file tạm cho video
-        driving_path = "driving_video.mp4"
+        driving_path = f"driving_video_{int(np.random.rand() * 10000)}.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")
+        # Kiểm tra nếu đã chọn sử dụng video mẫu
+        if driving_video_file is None:
+            # Tải và sử dụng video mẫu
+            sample_path = "sample_driving.mp4"
+            if not os.path.exists(sample_path) or os.path.getsize(sample_path) < 10000:
+                print("Đang tải video mẫu...")
+                os.system("wget -O sample_driving.mp4 https://github.com/AliaksandrSiarohin/first-order-model/raw/master/driving.mp4")
+            driving_path = sample_path
         else:
-            # Ghi dữ liệu nhị phân vào file
-            with open(driving_path, 'wb') as f:
-                f.write(driving_video_file)
+            # Xử lý video được tải lên
+            if isinstance(driving_video_file, str):
+                # Nếu là đường dẫn, copy file
+                if os.path.exists(driving_video_file):
+                    import shutil
+                    shutil.copyfile(driving_video_file, driving_path)
+                else:
+                    return None, f"Không tìm thấy file video tại đường dẫn: {driving_video_file}"
+            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,6 +279,13 @@ def animate_fomm(source_image, driving_video_file, relative=True, adapt_scale=Tr
             adapt_movement_scale=adapt_scale
         )
         
+        # Xóa file tạm nếu cần
+        if os.path.exists(source_path) and source_path != "source_image.jpg":
+            os.remove(source_path)
+        
+        if os.path.exists(driving_path) and driving_path != "sample_driving.mp4" and driving_path != "driving_video.mp4":
+            os.remove(driving_path)
+        
         return result_path, "Video được tạo thành công!"
     except Exception as e:
         import traceback
@@ -229,16 +300,16 @@ with gr.Blocks(title="First Order Motion Model - Tạo video người chuyển
         with gr.Column():
             source_image = gr.Image(type="pil", label="Tải lên ảnh nguồn")
             
+            # Thêm tùy chọn sử dụng video mẫu
+            use_sample = gr.Checkbox(label="Sử dụng video mẫu có sẵn", value=True)
+            
             # 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)")
+            driving_video_file = gr.File(label="Tải lên video tham chiếu (.mp4)", visible=False)
             
             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():
@@ -247,36 +318,32 @@ with gr.Blocks(title="First Order Motion Model - Tạo video người chuyển
     
     # 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)
+        return gr.update(visible=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):
+    def process_inputs(source_img, use_sample_vid, driving_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)
+            return animate_fomm(source_img, None, rel, adapt)
         else:
             return animate_fomm(source_img, driving_vid, rel, adapt)
     
     submit_btn.click(
         fn=process_inputs,
-        inputs=[source_image, driving_video_file, use_sample, relative, adapt_scale],
+        inputs=[source_image, use_sample, driving_video_file, 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. 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("2. Chọn sử dụng video mẫu có sẵn hoặc tải lên video tham chiếu của riêng bạn")
+    gr.Markdown("3. 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)")
     gr.Markdown("- Đối tượng nên ở vị trí tương tự trong cả ảnh nguồn và khung đầu tiên của video tham chiếu")
     gr.Markdown("- Quá trình tạo video có thể mất vài phút")
+    gr.Markdown("- Nếu gặp vấn đề với việc tải lên video, hãy sử dụng video mẫu có sẵn")
 
 demo.launch()