jhj0517 commited on
Commit
e7f3a8b
·
1 Parent(s): feac218

remove download models

Browse files
Files changed (2) hide show
  1. musepose_inference.py +4 -19
  2. pose_align.py +0 -15
musepose_inference.py CHANGED
@@ -23,15 +23,15 @@ from musepose.utils.util import get_fps, read_frames, save_videos_grid
23
  class MusePoseInference:
24
  def __init__(self):
25
  self.image_gen_model_paths = {
26
- "pretrained_base_model": "lambdalabs/sd-image-variations-diffusers/unet",
27
- "pretrained_vae": "stabilityai/sd-vae-ft-mse",
28
- "image_encoder": "lambdalabs/sd-image-variations-diffusers/image_encoder",
29
  }
30
  self.musepose_model_paths = {
31
  "denoising_unet": os.path.join("pretrained_weights", "MusePose", "denoising_unet.pth"),
32
  "reference_unet": os.path.join("pretrained_weights", "MusePose", "reference_unet.pth"),
33
  "pose_guider": os.path.join("pretrained_weights", "MusePose", "pose_guider.pth"),
34
- "motion_module": os.path.join("pretrained_weights", "MusePose", "pose_guider.pth"),
35
  }
36
  self.inference_config_path = os.path.join("configs", "inference_v2.yaml")
37
  self.vae = None
@@ -43,7 +43,6 @@ class MusePoseInference:
43
  self.output_dir = os.path.join("assets", "videos")
44
  if not os.path.exists(self.output_dir):
45
  os.makedirs(self.output_dir)
46
- self.download_models()
47
 
48
  def infer_musepose(
49
  self,
@@ -217,20 +216,6 @@ class MusePoseInference:
217
  self.release_vram()
218
  return output_path, output_path_demo
219
 
220
- def download_models(self):
221
- repo_id = 'jhj0517/MusePose'
222
- local_model_dir = os.path.join("pretrained_weights")
223
- for name, file_path in self.musepose_model_paths.items():
224
- local_dir, filename = os.path.dirname(file_path), os.path.basename(file_path)
225
- if not os.path.exists(local_dir):
226
- os.makedirs(local_dir)
227
-
228
- remote_filepath = f"MusePose/{filename}"
229
- if not os.path.exists(file_path):
230
- hf_hub_download(repo_id=repo_id, filename=remote_filepath,
231
- local_dir=local_model_dir,
232
- local_dir_use_symlinks=False)
233
-
234
  def release_vram(self):
235
  models = [
236
  'vae', 'reference_unet', 'denoising_unet',
 
23
  class MusePoseInference:
24
  def __init__(self):
25
  self.image_gen_model_paths = {
26
+ "pretrained_base_model": os.path.join("pretrained_weights", "sd-image-variations-diffusers"),
27
+ "pretrained_vae": os.path.join("pretrained_weights", "sd-vae-ft-mse"),
28
+ "image_encoder": os.path.join("pretrained_weights", "image_encoder"),
29
  }
30
  self.musepose_model_paths = {
31
  "denoising_unet": os.path.join("pretrained_weights", "MusePose", "denoising_unet.pth"),
32
  "reference_unet": os.path.join("pretrained_weights", "MusePose", "reference_unet.pth"),
33
  "pose_guider": os.path.join("pretrained_weights", "MusePose", "pose_guider.pth"),
34
+ "motion_module": os.path.join("pretrained_weights", "MusePose", "motion_module.pth"),
35
  }
36
  self.inference_config_path = os.path.join("configs", "inference_v2.yaml")
37
  self.vae = None
 
43
  self.output_dir = os.path.join("assets", "videos")
44
  if not os.path.exists(self.output_dir):
45
  os.makedirs(self.output_dir)
 
46
 
47
  def infer_musepose(
48
  self,
 
216
  self.release_vram()
217
  return output_path, output_path_demo
218
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
219
  def release_vram(self):
220
  models = [
221
  'vae', 'reference_unet', 'denoising_unet',
pose_align.py CHANGED
@@ -33,7 +33,6 @@ class PoseAlignmentInference:
33
  self.output_dir = os.path.join("assets", "videos")
34
  if not os.path.exists(self.output_dir):
35
  os.makedirs(self.output_dir)
36
- self.download_models()
37
 
38
  def align_pose(
39
  self,
@@ -312,20 +311,6 @@ class PoseAlignmentInference:
312
  self.release_vram()
313
  return outfn_align_pose_video, outfn
314
 
315
- def download_models(self):
316
- repo_id = 'jhj0517/MusePose'
317
- local_model_dir = os.path.join("pretrained_weights")
318
- for name, file_path in self.model_paths.items():
319
- local_dir, filename = os.path.dirname(file_path), os.path.basename(file_path)
320
- if not os.path.exists(local_dir):
321
- os.makedirs(local_dir)
322
-
323
- remote_filepath = f"dwpose/{filename}"
324
- if not os.path.exists(file_path):
325
- hf_hub_download(repo_id=repo_id, filename=remote_filepath,
326
- local_dir=local_model_dir,
327
- local_dir_use_symlinks=False)
328
-
329
  def release_vram(self):
330
  if self.detector is not None:
331
  del self.detector
 
33
  self.output_dir = os.path.join("assets", "videos")
34
  if not os.path.exists(self.output_dir):
35
  os.makedirs(self.output_dir)
 
36
 
37
  def align_pose(
38
  self,
 
311
  self.release_vram()
312
  return outfn_align_pose_video, outfn
313
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
314
  def release_vram(self):
315
  if self.detector is not None:
316
  del self.detector