Upload folder using huggingface_hub

animatediff/models/unet.py

@@ -456,12 +456,15 @@ class UNet3DConditionModel(ModelMixin, ConfigMixin):
         return UNet3DConditionOutput(sample=sample)
 
     @classmethod
-    def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, unet_additional_kwargs=None):
+    def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, unet_additional_kwargs=None, config_path=None):
         if subfolder is not None:
             pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
         print(f"loaded temporal unet's pretrained weights from {pretrained_model_path} ...")
-
         config_file = os.path.join(pretrained_model_path, 'config.json')
+        
+        if config_path is not None:
+            config_file = config_path
+            
         if not os.path.isfile(config_file):
             raise RuntimeError(f"{config_file} does not exist")
         with open(config_file, "r") as f:
@@ -482,7 +485,11 @@ class UNet3DConditionModel(ModelMixin, ConfigMixin):
 
         from diffusers.utils import WEIGHTS_NAME
         model = cls.from_config(config, **unet_additional_kwargs)
-        model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+        
+        if config_path is None:
+            model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+        else:
+            model_file = pretrained_model_path
         if not os.path.isfile(model_file):
             raise RuntimeError(f"{model_file} does not exist")
         state_dict = torch.load(model_file, map_location="cpu")

handler.py

@@ -3,7 +3,7 @@
 from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler
 from transformers import CLIPTextModel, CLIPTokenizer
 from omegaconf import OmegaConf
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, try_to_load_from_cache
 
 import os
 
@@ -21,8 +21,9 @@ from animatediff.utils.util import load_weights
 class EndpointHandler():
     def __init__(self, model_path: str = "bluestarburst/AnimateDiff-SceneFusion"):
         
-        inference_config_path = "configs/inference/inference-v3.yaml"
-        hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="configs/inference/inference-v3.yaml")
+        # inference_config_path = "configs/inference/inference-v3.yaml"
+        inference_config_path = hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="configs/inference/inference-v3.yaml")
+        print(inference_config_path)
         
         inference_config = OmegaConf.load(inference_config_path)
         
@@ -33,13 +34,12 @@ class EndpointHandler():
         text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="models/StableDiffusion/text_encoder")
         vae          = AutoencoderKL.from_pretrained(model_path, subfolder="models/StableDiffusion/vae")         
         
-        if not os.path.isfile("models/StableDiffusion/unet/diffusion_pytorch_model.bin"):
-            hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/StableDiffusion/unet/config.json")
-            hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/StableDiffusion/unet/diffusion_pytorch_model.bin")
+        unet_model_path = hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/StableDiffusion/unet/diffusion_pytorch_model.bin")
+        unet_config_path = hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/StableDiffusion/unet/config.json")        
 
-        unet_model_path = "models/StableDiffusion/unet"
+        print(unet_model_path)
 
-        unet         = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path=unet_model_path, unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
+        unet         = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path=unet_model_path, unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs), config_path=unet_config_path)
 
         if is_xformers_available(): unet.enable_xformers_memory_efficient_attention()
         else: assert False