Spaces:
Runtime error
Runtime error
| import logging | |
| from functools import wraps | |
| from pathlib import Path | |
| from typing import Optional, TypeVar | |
| from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline | |
| from huggingface_hub import hf_hub_download | |
| from torch import nn | |
| from animatediff import HF_HUB_CACHE, HF_MODULE_REPO, get_dir | |
| from animatediff.settings import CKPT_EXTENSIONS | |
| from animatediff.utils.huggingface import get_hf_pipeline, get_hf_pipeline_sdxl | |
| from animatediff.utils.util import path_from_cwd | |
| logger = logging.getLogger(__name__) | |
| data_dir = get_dir("data") | |
| checkpoint_dir = data_dir.joinpath("models/sd") | |
| pipeline_dir = data_dir.joinpath("models/huggingface") | |
| # for the nop_train() monkeypatch | |
| T = TypeVar("T", bound=nn.Module) | |
| def nop_train(self: T, mode: bool = True) -> T: | |
| """No-op for monkeypatching train() call to prevent unfreezing module""" | |
| return self | |
| def get_base_model(model_name_or_path: str, local_dir: Path, force: bool = False, is_sdxl:bool=False) -> Path: | |
| model_name_or_path = Path(model_name_or_path) | |
| model_save_dir = local_dir.joinpath(str(model_name_or_path).split("/")[-1]).resolve() | |
| model_is_repo_id = False if model_name_or_path.joinpath("model_index.json").exists() else True | |
| # if we have a HF repo ID, download it | |
| if model_is_repo_id: | |
| logger.debug("Base model is a HuggingFace repo ID") | |
| if model_save_dir.joinpath("model_index.json").exists(): | |
| logger.debug(f"Base model already downloaded to: {path_from_cwd(model_save_dir)}") | |
| else: | |
| logger.info(f"Downloading base model from {model_name_or_path}...") | |
| if is_sdxl: | |
| _ = get_hf_pipeline_sdxl(model_name_or_path, model_save_dir, save=True, force_download=force) | |
| else: | |
| _ = get_hf_pipeline(model_name_or_path, model_save_dir, save=True, force_download=force) | |
| model_name_or_path = model_save_dir | |
| return Path(model_name_or_path) | |
| def fix_checkpoint_if_needed(checkpoint: Path, debug:bool): | |
| def dump(loaded): | |
| for a in loaded: | |
| logger.info(f"{a} {loaded[a].shape}") | |
| if debug: | |
| from safetensors.torch import load_file, save_file | |
| loaded = load_file(checkpoint, "cpu") | |
| dump(loaded) | |
| return | |
| try: | |
| pipeline = StableDiffusionPipeline.from_single_file( | |
| pretrained_model_link_or_path=str(checkpoint.absolute()), | |
| local_files_only=False, | |
| load_safety_checker=False, | |
| ) | |
| logger.info("This file works fine.") | |
| return | |
| except: | |
| from safetensors.torch import load_file, save_file | |
| loaded = load_file(checkpoint, "cpu") | |
| convert_table_bias={ | |
| "first_stage_model.decoder.mid.attn_1.to_k.bias":"first_stage_model.decoder.mid.attn_1.k.bias", | |
| "first_stage_model.decoder.mid.attn_1.to_out.0.bias":"first_stage_model.decoder.mid.attn_1.proj_out.bias", | |
| "first_stage_model.decoder.mid.attn_1.to_q.bias":"first_stage_model.decoder.mid.attn_1.q.bias", | |
| "first_stage_model.decoder.mid.attn_1.to_v.bias":"first_stage_model.decoder.mid.attn_1.v.bias", | |
| "first_stage_model.encoder.mid.attn_1.to_k.bias":"first_stage_model.encoder.mid.attn_1.k.bias", | |
| "first_stage_model.encoder.mid.attn_1.to_out.0.bias":"first_stage_model.encoder.mid.attn_1.proj_out.bias", | |
| "first_stage_model.encoder.mid.attn_1.to_q.bias":"first_stage_model.encoder.mid.attn_1.q.bias", | |
| "first_stage_model.encoder.mid.attn_1.to_v.bias":"first_stage_model.encoder.mid.attn_1.v.bias", | |
| } | |
| convert_table_weight={ | |
| "first_stage_model.decoder.mid.attn_1.to_k.weight":"first_stage_model.decoder.mid.attn_1.k.weight", | |
| "first_stage_model.decoder.mid.attn_1.to_out.0.weight":"first_stage_model.decoder.mid.attn_1.proj_out.weight", | |
| "first_stage_model.decoder.mid.attn_1.to_q.weight":"first_stage_model.decoder.mid.attn_1.q.weight", | |
| "first_stage_model.decoder.mid.attn_1.to_v.weight":"first_stage_model.decoder.mid.attn_1.v.weight", | |
| "first_stage_model.encoder.mid.attn_1.to_k.weight":"first_stage_model.encoder.mid.attn_1.k.weight", | |
| "first_stage_model.encoder.mid.attn_1.to_out.0.weight":"first_stage_model.encoder.mid.attn_1.proj_out.weight", | |
| "first_stage_model.encoder.mid.attn_1.to_q.weight":"first_stage_model.encoder.mid.attn_1.q.weight", | |
| "first_stage_model.encoder.mid.attn_1.to_v.weight":"first_stage_model.encoder.mid.attn_1.v.weight", | |
| } | |
| for a in list(loaded.keys()): | |
| if a in convert_table_bias: | |
| new_key = convert_table_bias[a] | |
| loaded[new_key] = loaded.pop(a) | |
| elif a in convert_table_weight: | |
| new_key = convert_table_weight[a] | |
| item = loaded.pop(a) | |
| if len(item.shape) == 2: | |
| item = item.unsqueeze(dim=-1).unsqueeze(dim=-1) | |
| loaded[new_key] = item | |
| new_path = str(checkpoint.parent / checkpoint.stem) + "_fixed"+checkpoint.suffix | |
| logger.info(f"Saving file to {new_path}") | |
| save_file(loaded, Path(new_path)) | |
| def checkpoint_to_pipeline( | |
| checkpoint: Path, | |
| target_dir: Optional[Path] = None, | |
| save: bool = True, | |
| ) -> StableDiffusionPipeline: | |
| logger.debug(f"Converting checkpoint {path_from_cwd(checkpoint)}") | |
| if target_dir is None: | |
| target_dir = pipeline_dir.joinpath(checkpoint.stem) | |
| pipeline = StableDiffusionPipeline.from_single_file( | |
| pretrained_model_link_or_path=str(checkpoint.absolute()), | |
| local_files_only=False, | |
| load_safety_checker=False, | |
| ) | |
| if save: | |
| target_dir.mkdir(parents=True, exist_ok=True) | |
| logger.info(f"Saving pipeline to {path_from_cwd(target_dir)}") | |
| pipeline.save_pretrained(target_dir, safe_serialization=True) | |
| return pipeline, target_dir | |
| def checkpoint_to_pipeline_sdxl( | |
| checkpoint: Path, | |
| target_dir: Optional[Path] = None, | |
| save: bool = True, | |
| ) -> StableDiffusionXLPipeline: | |
| logger.debug(f"Converting checkpoint {path_from_cwd(checkpoint)}") | |
| if target_dir is None: | |
| target_dir = pipeline_dir.joinpath(checkpoint.stem) | |
| pipeline = StableDiffusionXLPipeline.from_single_file( | |
| pretrained_model_link_or_path=str(checkpoint.absolute()), | |
| local_files_only=False, | |
| load_safety_checker=False, | |
| ) | |
| if save: | |
| target_dir.mkdir(parents=True, exist_ok=True) | |
| logger.info(f"Saving pipeline to {path_from_cwd(target_dir)}") | |
| pipeline.save_pretrained(target_dir, safe_serialization=True) | |
| return pipeline, target_dir | |
| def get_checkpoint_weights(checkpoint: Path): | |
| temp_pipeline: StableDiffusionPipeline | |
| temp_pipeline, _ = checkpoint_to_pipeline(checkpoint, save=False) | |
| unet_state_dict = temp_pipeline.unet.state_dict() | |
| tenc_state_dict = temp_pipeline.text_encoder.state_dict() | |
| vae_state_dict = temp_pipeline.vae.state_dict() | |
| return unet_state_dict, tenc_state_dict, vae_state_dict | |
| def get_checkpoint_weights_sdxl(checkpoint: Path): | |
| temp_pipeline: StableDiffusionXLPipeline | |
| temp_pipeline, _ = checkpoint_to_pipeline_sdxl(checkpoint, save=False) | |
| unet_state_dict = temp_pipeline.unet.state_dict() | |
| tenc_state_dict = temp_pipeline.text_encoder.state_dict() | |
| tenc2_state_dict = temp_pipeline.text_encoder_2.state_dict() | |
| vae_state_dict = temp_pipeline.vae.state_dict() | |
| return unet_state_dict, tenc_state_dict, tenc2_state_dict, vae_state_dict | |
| def ensure_motion_modules( | |
| repo_id: str = HF_MODULE_REPO, | |
| fp16: bool = False, | |
| force: bool = False, | |
| ): | |
| """Retrieve the motion modules from HuggingFace Hub.""" | |
| module_files = ["mm_sd_v14.safetensors", "mm_sd_v15.safetensors"] | |
| module_dir = get_dir("data/models/motion-module") | |
| for file in module_files: | |
| target_path = module_dir.joinpath(file) | |
| if fp16: | |
| target_path = target_path.with_suffix(".fp16.safetensors") | |
| if target_path.exists() and force is not True: | |
| logger.debug(f"File {path_from_cwd(target_path)} already exists, skipping download") | |
| else: | |
| result = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=target_path.name, | |
| cache_dir=HF_HUB_CACHE, | |
| local_dir=module_dir, | |
| local_dir_use_symlinks=False, | |
| resume_download=True, | |
| ) | |
| logger.debug(f"Downloaded {path_from_cwd(result)}") | |