Spaces:
Sleeping
Sleeping
jhj0517
commited on
Commit
·
e7f3a8b
1
Parent(s):
feac218
remove download models
Browse files- musepose_inference.py +4 -19
- 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": "
|
27 |
-
"pretrained_vae": "
|
28 |
-
"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", "
|
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
|