Upload folder using huggingface_hub

handler.py

@@ -5,6 +5,7 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from omegaconf import OmegaConf
 from huggingface_hub import hf_hub_download
 
+import os
 
 from diffusers.utils.import_utils import is_xformers_available
 from typing import Any
@@ -20,8 +21,8 @@ from animatediff.utils.util import load_weights
 class EndpointHandler():
     def __init__(self, model_path: str = "bluestarburst/AnimateDiff-SceneFusion"):
         
-        inference_config_path = "configs/inference-v3.yaml"
-        hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="configs/inference/inference-v3.yaml", output_dir="configs")
+        inference_config_path = "configs/inference/inference-v3.yaml"
+        hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="configs/inference/inference-v3.yaml")
         
         inference_config = OmegaConf.load(inference_config_path)
         
@@ -30,21 +31,28 @@ class EndpointHandler():
         ### >>> create validation pipeline >>> ###
         tokenizer    = CLIPTokenizer.from_pretrained(model_path, subfolder="models/StableDiffusion/tokenizer")
         text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="models/StableDiffusion/text_encoder")
-        vae          = AutoencoderKL.from_pretrained(model_path, subfolder="models/StableDiffusion/vae")            
-        unet         = UNet3DConditionModel.from_pretrained_2d(model_path, subfolder="models/StableDiffusion/unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
+        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 = "models/StableDiffusion/unet"
+
+        unet         = UNet3DConditionModel.from_pretrained_2d(pretrained_model_path=unet_model_path, unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
 
         if is_xformers_available(): unet.enable_xformers_memory_efficient_attention()
         else: assert False
 
         self.pipeline = AnimationPipeline(
             vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
-            scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config['noise_scheduler_kwargs']['DDIMScheduler']'])),
+            scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs.DDIMScheduler))
         ).to("cuda")
         
         # huggingface download motion module from bluestarburst/AnimateDiff-SceneFusion/models/Motion_Module/mm_sd_v15.ckpt
 
-        motion_module = "models/Motion_Module/mm_sd_v15.ckpt"
-        hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/Motion_Module/mm_sd_v15.ckpt", output_dir="models/Motion_Module")
+        motion_module = "models/MotionModule/mm_sd_v15.ckpt"
+        hf_hub_download(repo_id="bluestarburst/AnimateDiff-SceneFusion", filename="models/Motion_Module/mm_sd_v15.ckpt")
         
 
         self.pipeline = load_weights(
@@ -97,4 +105,4 @@ class EndpointHandler():
 # This function will be called during inference time.
 
 
-# new_handler = EndpointHandler()
+new_handler = EndpointHandler()