swapping models

download_weights.py

@@ -0,0 +1,111 @@
+import os
+from pathlib import Path, PurePosixPath
+
+from huggingface_hub import hf_hub_download
+
+
+def prepare_base_model():
+    print(f'Preparing base stable-diffusion-v1-5 weights...')
+    local_dir = "./pretrained_weights/stable-diffusion-v1-5"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in ["unet/config.json", "unet/diffusion_pytorch_model.bin"]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+        hf_hub_download(
+            repo_id="runwayml/stable-diffusion-v1-5",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_image_encoder():
+    print(f"Preparing image encoder weights...")
+    local_dir = "./pretrained_weights"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in ["image_encoder/config.json", "image_encoder/pytorch_model.bin"]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+        hf_hub_download(
+            repo_id="lambdalabs/sd-image-variations-diffusers",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_dwpose():
+    print(f"Preparing DWPose weights...")
+    local_dir = "./pretrained_weights/DWPose"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "dw-ll_ucoco_384.onnx",
+        "yolox_l.onnx",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="yzd-v/DWPose",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_vae():
+    print(f"Preparing vae weights...")
+    local_dir = "./pretrained_weights/sd-vae-ft-mse"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "config.json",
+        "diffusion_pytorch_model.bin",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="stabilityai/sd-vae-ft-mse",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+
+def prepare_anyone():
+    print(f"Preparing AnimateAnyone weights...")
+    local_dir = "./pretrained_weights"
+    os.makedirs(local_dir, exist_ok=True)
+    for hub_file in [
+        "denoising_unet.pth",
+        "motion_module.pth",
+        "pose_guider.pth",
+        "reference_unet.pth",
+    ]:
+        path = Path(hub_file)
+        saved_path = local_dir / path
+        if os.path.exists(saved_path):
+            continue
+
+        hf_hub_download(
+            repo_id="patrolli/AnimateAnyone",
+            subfolder=PurePosixPath(path.parent),
+            filename=PurePosixPath(path.name),
+            local_dir=local_dir,
+        )
+
+if __name__ == '__main__':
+    prepare_base_model()
+    prepare_image_encoder()
+    prepare_dwpose()
+    prepare_vae()
+    prepare_anyone()
+    

handler.py

@@ -50,8 +50,7 @@ class EndpointHandler():
         pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
         print("model path is " + pretrained_base_model_path_unet)
         reference_unet = UNet2DConditionModel.from_pretrained(
-            pretrained_base_model_path_unet,
-            use_safetensors = True
+            pretrained_base_model_path_unet
         ).to(dtype=self.weight_dtype, device="cuda")
 
         inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')

pretrained_weights/sd-vae-ft-mse/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815
-size 334643276

pretrained_weights/stable-diffusion-v1-5/feature_extractor/preprocessor_config.json

@@ -1,20 +0,0 @@
-{
-  "crop_size": 224,
-  "do_center_crop": true,
-  "do_convert_rgb": true,
-  "do_normalize": true,
-  "do_resize": true,
-  "feature_extractor_type": "CLIPFeatureExtractor",
-  "image_mean": [
-    0.48145466,
-    0.4578275,
-    0.40821073
-  ],
-  "image_std": [
-    0.26862954,
-    0.26130258,
-    0.27577711
-  ],
-  "resample": 3,
-  "size": 224
-}

pretrained_weights/stable-diffusion-v1-5/model_index.json

@@ -1,32 +0,0 @@
-{
-  "_class_name": "StableDiffusionPipeline",
-  "_diffusers_version": "0.6.0",
-  "feature_extractor": [
-    "transformers",
-    "CLIPImageProcessor"
-  ],
-  "safety_checker": [
-    "stable_diffusion",
-    "StableDiffusionSafetyChecker"
-  ],
-  "scheduler": [
-    "diffusers",
-    "PNDMScheduler"
-  ],
-  "text_encoder": [
-    "transformers",
-    "CLIPTextModel"
-  ],
-  "tokenizer": [
-    "transformers",
-    "CLIPTokenizer"
-  ],
-  "unet": [
-    "diffusers",
-    "UNet2DConditionModel"
-  ],
-  "vae": [
-    "diffusers",
-    "AutoencoderKL"
-  ]
-}

pretrained_weights/stable-diffusion-v1-5/v1-inference.yaml

@@ -1,70 +0,0 @@
-model:
-  base_learning_rate: 1.0e-04
-  target: ldm.models.diffusion.ddpm.LatentDiffusion
-  params:
-    linear_start: 0.00085
-    linear_end: 0.0120
-    num_timesteps_cond: 1
-    log_every_t: 200
-    timesteps: 1000
-    first_stage_key: "jpg"
-    cond_stage_key: "txt"
-    image_size: 64
-    channels: 4
-    cond_stage_trainable: false   # Note: different from the one we trained before
-    conditioning_key: crossattn
-    monitor: val/loss_simple_ema
-    scale_factor: 0.18215
-    use_ema: False
-
-    scheduler_config: # 10000 warmup steps
-      target: ldm.lr_scheduler.LambdaLinearScheduler
-      params:
-        warm_up_steps: [ 10000 ]
-        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
-        f_start: [ 1.e-6 ]
-        f_max: [ 1. ]
-        f_min: [ 1. ]
-
-    unet_config:
-      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
-      params:
-        image_size: 32 # unused
-        in_channels: 4
-        out_channels: 4
-        model_channels: 320
-        attention_resolutions: [ 4, 2, 1 ]
-        num_res_blocks: 2
-        channel_mult: [ 1, 2, 4, 4 ]
-        num_heads: 8
-        use_spatial_transformer: True
-        transformer_depth: 1
-        context_dim: 768
-        use_checkpoint: True
-        legacy: False
-
-    first_stage_config:
-      target: ldm.models.autoencoder.AutoencoderKL
-      params:
-        embed_dim: 4
-        monitor: val/rec_loss
-        ddconfig:
-          double_z: true
-          z_channels: 4
-          resolution: 256
-          in_channels: 3
-          out_ch: 3
-          ch: 128
-          ch_mult:
-          - 1
-          - 2
-          - 4
-          - 4
-          num_res_blocks: 2
-          attn_resolutions: []
-          dropout: 0.0
-        lossconfig:
-          target: torch.nn.Identity
-
-    cond_stage_config:
-      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder

requirements.txt

@@ -24,19 +24,18 @@ gradio==3.41.2
 onnxruntime-coreml==1.13.1; python_version == '3.9' and sys_platform == 'darwin' and platform_machine != 'arm64'
 transformers==4.41.1
 
-# Stuff huggingface is complaining about
-einops==0.4.1
-av==11.0.0
-
 # Add additional dependencies
 diffusers==0.24.0
 omegaconf==2.2.3
-safetensors==0.4.1
 
 # Face swap related dependencies
 facenet-pytorch==2.5.2
 dlib==19.22.0
 
+# Stuff huggingface is complaining about
+einops==0.4.1
+av==11.0.0
+
 # Additional dependencies from the first list not present in the second list
 accelerate==0.21.0
 clip @ https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip#sha256=b5842c25da441d6c581b53a5c60e0c2127ebafe0f746f8e15561a006c6c3be6a