alvinichi commited on
Commit
5b61f28
·
1 Parent(s): 6be96f1
Files changed (1) hide show
  1. app.py +149 -76
app.py CHANGED
@@ -9,15 +9,16 @@ from skimage import img_as_ubyte
9
  import gradio as gr
10
  from PIL import Image
11
  import tempfile
 
 
12
 
13
- # Đảm bảo cài đặt các thư viện cần thiết trước khi import
14
  subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-learn"])
15
  subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-image==0.19.3"])
16
  subprocess.check_call([sys.executable, "-m", "pip", "install", "face-alignment==1.3.5"])
17
  subprocess.check_call([sys.executable, "-m", "pip", "install", "PyYAML==5.3.1"])
18
  subprocess.check_call([sys.executable, "-m", "pip", "install", "imageio-ffmpeg==0.4.5"])
19
- subprocess.check_call([sys.executable, "-m", "pip", "install", "gdown"])
20
- subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
21
 
22
  # Cài đặt ffmpeg trong môi trường Ubuntu
23
  os.system("apt-get update && apt-get install -y ffmpeg")
@@ -95,77 +96,100 @@ def normalize_kp(kp_source, kp_driving, kp_driving_initial,
95
  # Import hàm load_checkpoints từ file helper
96
  from load_helper import load_checkpoints, normalize_kp
97
 
98
- # Tải mô hình pre-trained với phương pháp cải tiến
99
  def download_model():
100
- try:
101
- # Thử phương pháp sử dụng gdown trước
102
- model_path = 'checkpoints/vox-cpk.pth.tar'
103
- if not os.path.exists('checkpoints'):
104
- os.makedirs('checkpoints', exist_ok=True)
105
-
106
- # Kiểm tra xem file đã tồn tại và đủ lớn chưa
107
- if not os.path.exists(model_path) or os.path.getsize(model_path) < 1000000:
108
- print("Đang tải mô hình từ Google Drive...")
109
- import gdown
110
- file_id = '1PyQJmkdCsAkOYwUyaj_l-l0as-iLDgeH'
111
- gdown.download(f"https://drive.google.com/uc?id={file_id}", model_path, quiet=False)
112
-
113
- config_path = 'first_order_model/config/vox-256.yaml'
114
- if not os.path.exists('first_order_model/config'):
115
- os.makedirs('first_order_model/config', exist_ok=True)
116
-
117
- if not os.path.exists(config_path) or os.path.getsize(config_path) < 1000:
118
- print("Đang tải file cấu hình từ Google Drive...")
119
- import gdown
120
- file_id = '1pZUMNRjkBiuBEM68oj9nskuWgJR-5QMn'
121
- gdown.download(f"https://drive.google.com/uc?id={file_id}", config_path, quiet=False)
122
-
123
- # Kiểm tra lại kích thước file
124
- 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:
125
- return config_path, model_path
126
-
127
- # Nếu gdown không thành công, thử phương pháp sử dụng Hugging Face
128
- raise Exception("Tải mô hình từ Google Drive không thành công. Chuyển sang sử dụng Hugging Face...")
129
 
130
- except Exception as e:
131
- print(f"Lỗi khi tải từ Google Drive: {str(e)}")
132
- return download_model_alternative()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
133
 
134
- def download_model_alternative():
135
- try:
136
- print("Đang tải hình từ Hugging Face...")
137
- from huggingface_hub import hf_hub_download
138
-
139
- model_path = hf_hub_download(
140
- repo_id="ycyunwei/first-order-motion-model",
141
- filename="vox-cpk.pth.tar"
142
- )
143
-
144
- config_path = hf_hub_download(
145
- repo_id="ycyunwei/first-order-motion-model",
146
- filename="vox-256.yaml"
147
- )
148
-
149
- return config_path, model_path
150
- except Exception as e:
151
- print(f"Lỗi khi tải từ Hugging Face: {str(e)}")
152
- # Thử phương pháp cuối cùng - sử dụng direct link
153
- model_path = 'checkpoints/vox-cpk.pth.tar'
154
- if not os.path.exists('checkpoints'):
155
- os.makedirs('checkpoints', exist_ok=True)
156
-
157
- os.system(f'wget -O {model_path} https://github.com/AliaksandrSiarohin/first-order-model/releases/download/v1.0.0/vox-cpk.pth.tar')
158
-
159
- config_path = 'first_order_model/config/vox-256.yaml'
160
- if not os.path.exists('first_order_model/config'):
161
- os.makedirs('first_order_model/config', exist_ok=True)
162
-
163
- os.system(f'wget -O {config_path} https://raw.githubusercontent.com/AliaksandrSiarohin/first-order-model/master/config/vox-256.yaml')
164
-
165
- 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:
166
- return config_path, model_path
167
- else:
168
- 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.")
169
 
170
  # Hàm tạo animation
171
  def make_animation(source_image, driving_video, relative=True, adapt_movement_scale=True):
@@ -233,6 +257,59 @@ def make_animation(source_image, driving_video, relative=True, adapt_movement_sc
233
 
234
  return output_path
235
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
236
  # Định nghĩa giao diện Gradio
237
  def animate_fomm(source_image, driving_video_file, relative=True, adapt_scale=True):
238
  if source_image is None:
@@ -252,11 +329,7 @@ def animate_fomm(source_image, driving_video_file, relative=True, adapt_scale=Tr
252
  # Kiểm tra nếu đã chọn sử dụng video mẫu
253
  if driving_video_file is None:
254
  # Tải và sử dụng video mẫu
255
- sample_path = "sample_driving.mp4"
256
- if not os.path.exists(sample_path) or os.path.getsize(sample_path) < 10000:
257
- print("Đang tải video mẫu...")
258
- os.system("wget -O sample_driving.mp4 https://github.com/AliaksandrSiarohin/first-order-model/raw/master/driving.mp4")
259
- driving_path = sample_path
260
  else:
261
  # Xử lý video được tải lên
262
  if isinstance(driving_video_file, str):
 
9
  import gradio as gr
10
  from PIL import Image
11
  import tempfile
12
+ import requests
13
+ from io import BytesIO
14
 
15
+ # Đảm bảo cài đặt các thư viện cần thiết
16
  subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-learn"])
17
  subprocess.check_call([sys.executable, "-m", "pip", "install", "scikit-image==0.19.3"])
18
  subprocess.check_call([sys.executable, "-m", "pip", "install", "face-alignment==1.3.5"])
19
  subprocess.check_call([sys.executable, "-m", "pip", "install", "PyYAML==5.3.1"])
20
  subprocess.check_call([sys.executable, "-m", "pip", "install", "imageio-ffmpeg==0.4.5"])
21
+ subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])
 
22
 
23
  # Cài đặt ffmpeg trong môi trường Ubuntu
24
  os.system("apt-get update && apt-get install -y ffmpeg")
 
96
  # Import hàm load_checkpoints từ file helper
97
  from load_helper import load_checkpoints, normalize_kp
98
 
99
+ # Tải mô hình từ GitHub hoặc mirrors của first-order-model
100
  def download_model():
101
+ # URLs trực tiếp từ sources khác
102
+ checkpoint_urls = [
103
+ "https://github.com/AliaksandrSiarohin/first-order-model/releases/download/v1.0.0/vox-cpk.pth.tar",
104
+ "https://raw.githubusercontent.com/jiupinjia/stylized-neural-painting/main/checkpoints/vox-cpk.pth.tar",
105
+ "https://github.com/snap-research/articulated-animation/raw/master/checkpoints/vox.pth.tar"
106
+ ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
107
 
108
+ config_urls = [
109
+ "https://raw.githubusercontent.com/AliaksandrSiarohin/first-order-model/master/config/vox-256.yaml",
110
+ "https://gist.githubusercontent.com/anonymous/raw/vox-256.yaml"
111
+ ]
112
+
113
+ # Tạo thư mục
114
+ model_path = 'checkpoints/vox-cpk.pth.tar'
115
+ if not os.path.exists('checkpoints'):
116
+ os.makedirs('checkpoints', exist_ok=True)
117
+
118
+ config_path = 'first_order_model/config/vox-256.yaml'
119
+ if not os.path.exists('first_order_model/config'):
120
+ os.makedirs('first_order_model/config', exist_ok=True)
121
+
122
+ # Tải model checkpoint
123
+ success = False
124
+ for url in checkpoint_urls:
125
+ try:
126
+ print(f"Đang thử tải mô hình từ: {url}")
127
+ response = requests.get(url, stream=True, timeout=30)
128
+ if response.status_code == 200:
129
+ with open(model_path, 'wb') as f:
130
+ for chunk in response.iter_content(chunk_size=8192):
131
+ f.write(chunk)
132
+
133
+ # Kiểm tra kích thước file (checkpoint mô hình thường > 100MB)
134
+ if os.path.getsize(model_path) > 100000000:
135
+ success = True
136
+ break
137
+ except Exception as e:
138
+ print(f"Lỗi khi tải từ {url}: {str(e)}")
139
+
140
+ if not success:
141
+ raise Exception("Không thể tải mô hình checkpoint từ bất kỳ nguồn nào")
142
+
143
+ # Tải file cấu h��nh
144
+ config_success = False
145
+ for url in config_urls:
146
+ try:
147
+ print(f"Đang thử tải file cấu hình từ: {url}")
148
+ response = requests.get(url, timeout=30)
149
+ if response.status_code == 200:
150
+ with open(config_path, 'wb') as f:
151
+ f.write(response.content)
152
+
153
+ if os.path.getsize(config_path) > 1000:
154
+ config_success = True
155
+ break
156
+ except Exception as e:
157
+ print(f"Lỗi khi tải cấu hình từ {url}: {str(e)}")
158
+
159
+ if not config_success:
160
+ # Tạo file cấu hình đơn giản nếu không tải được
161
+ create_simple_config(config_path)
162
+
163
+ return config_path, model_path
164
 
165
+ # Tạo file cấu hình đơn giản nếu không tải được
166
+ def create_simple_config(config_path):
167
+ with open(config_path, 'w') as f:
168
+ f.write("""
169
+ model_params:
170
+ common_params:
171
+ num_kp: 10
172
+ num_channels: 3
173
+ estimate_jacobian: true
174
+ kp_detector_params:
175
+ temperature: 0.1
176
+ block_expansion: 32
177
+ max_features: 1024
178
+ scale_factor: 0.25
179
+ num_blocks: 5
180
+ generator_params:
181
+ block_expansion: 64
182
+ max_features: 512
183
+ num_down_blocks: 2
184
+ num_bottleneck_blocks: 6
185
+ estimate_occlusion_map: true
186
+ dense_motion_params:
187
+ block_expansion: 64
188
+ max_features: 1024
189
+ num_blocks: 5
190
+ scale_factor: 0.25
191
+ """)
192
+ print("Đã tạo file cấu hình đơn giản")
 
 
 
 
 
 
 
193
 
194
  # Hàm tạo animation
195
  def make_animation(source_image, driving_video, relative=True, adapt_movement_scale=True):
 
257
 
258
  return output_path
259
 
260
+ # Tải video mẫu
261
+ def download_sample_video():
262
+ sample_urls = [
263
+ "https://github.com/AliaksandrSiarohin/first-order-model/raw/master/driving.mp4",
264
+ "https://raw.githubusercontent.com/jiupinjia/stylized-neural-painting/main/sample/driving.mp4"
265
+ ]
266
+
267
+ sample_path = "sample_driving.mp4"
268
+
269
+ for url in sample_urls:
270
+ try:
271
+ print(f"Đang thử tải video mẫu từ: {url}")
272
+ response = requests.get(url, timeout=30)
273
+ if response.status_code == 200:
274
+ with open(sample_path, 'wb') as f:
275
+ f.write(response.content)
276
+
277
+ if os.path.getsize(sample_path) > 10000: # Kiểm tra kích thước file
278
+ return sample_path
279
+ except Exception as e:
280
+ print(f"Lỗi khi tải video mẫu từ {url}: {str(e)}")
281
+
282
+ # Nếu không tải được, tạo video đơn giản
283
+ create_simple_video(sample_path)
284
+ return sample_path
285
+
286
+ # Tạo video đơn giản nếu không tải được video mẫu
287
+ def create_simple_video(output_path):
288
+ import cv2
289
+ out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 10, (256, 256))
290
+
291
+ # Tạo 100 khung hình với chuyển động đơn giản
292
+ for i in range(100):
293
+ frame = np.zeros((256, 256, 3), dtype=np.uint8)
294
+
295
+ # Vẽ khuôn mặt đơn giản chuyển động
296
+ x_center = 128 + int(np.sin(i/10) * 20)
297
+ y_center = 128 + int(np.cos(i/20) * 10)
298
+
299
+ # Vẽ khuôn mặt
300
+ cv2.circle(frame, (x_center, y_center), 60, (200, 200, 200), -1) # Mặt
301
+ cv2.circle(frame, (x_center - 20, y_center - 15), 10, (0, 0, 0), -1) # Mắt trái
302
+ cv2.circle(frame, (x_center + 20, y_center - 15), 10, (0, 0, 0), -1) # Mắt phải
303
+
304
+ # Vẽ miệng
305
+ mouth_y = y_center + 20 + int(np.sin(i/5) * 5)
306
+ cv2.ellipse(frame, (x_center, mouth_y), (20, 10), 0, 0, 180, (0, 0, 0), -1)
307
+
308
+ out.write(frame)
309
+
310
+ out.release()
311
+ print("Đã tạo video đơn giản")
312
+
313
  # Định nghĩa giao diện Gradio
314
  def animate_fomm(source_image, driving_video_file, relative=True, adapt_scale=True):
315
  if source_image is None:
 
329
  # Kiểm tra nếu đã chọn sử dụng video mẫu
330
  if driving_video_file is None:
331
  # Tải và sử dụng video mẫu
332
+ driving_path = download_sample_video()
 
 
 
 
333
  else:
334
  # Xử lý video được tải lên
335
  if isinstance(driving_video_file, str):