This commit is contained in:
Yuwei Guo
2023-07-09 23:25:46 +08:00
parent 5b702ae4e9
commit ebfd7b74f7
3 changed files with 96 additions and 84 deletions

View File

@@ -1,7 +1,9 @@
ToonYou: ToonYou:
base: "" base: ""
path: "models/DreamBooth_LoRA/toonyou_beta3.safetensors" path: "models/DreamBooth_LoRA/toonyou_beta3.safetensors"
motion_module: "models/Motion_Module/mm_sd_v14.ckpt" motion_module:
- "models/Motion_Module/mm_sd_v14.ckpt"
- "models/Motion_Module/mm_sd_v15.ckpt"
seed: [10788741199826055526, 6520604954829636163, 6519455744612555650, 16372571278361863751] seed: [10788741199826055526, 6520604954829636163, 6519455744612555650, 16372571278361863751]
steps: 25 steps: 25

View File

@@ -3,7 +3,8 @@ torch==1.12.1+cu113
torchvision==0.13.1+cu113 torchvision==0.13.1+cu113
diffusers[torch]==0.11.1 diffusers[torch]==0.11.1
transformers==4.25.1 transformers==4.25.1
imageio==2.27.0
gdown
einops einops
omegaconf omegaconf
safetensors safetensors
imageio==2.27.0

View File

@@ -37,94 +37,103 @@ def main(args):
config = OmegaConf.load(args.config) config = OmegaConf.load(args.config)
samples = [] samples = []
sample_idx = 0
for model_idx, (config_key, model_config) in enumerate(list(config.items())): for model_idx, (config_key, model_config) in enumerate(list(config.items())):
### >>> create validation pipeline >>> ###
tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_path, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_path, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder="vae")
unet = UNet3DConditionModel.from_pretrained_2d(args.pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
pipeline = AnimationPipeline( motion_modules = model_config.motion_module
vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, motion_modules = [motion_modules] if isinstance(motion_modules, str) else list(motion_modules)
scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs)), for motion_module in motion_modules:
).to("cuda")
# 1. unet ckpt ### >>> create validation pipeline >>> ###
# 1.1 motion module tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_path, subfolder="tokenizer")
motion_module_state_dict = torch.load(model_config.motion_module, map_location="cpu") text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_path, subfolder="text_encoder")
if "global_step" in motion_module_state_dict: func_args.update({"global_step": motion_module_state_dict["global_step"]}) vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder="vae")
missing, unexpected = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False) unet = UNet3DConditionModel.from_pretrained_2d(args.pretrained_model_path, subfolder="unet", unet_additional_kwargs=OmegaConf.to_container(inference_config.unet_additional_kwargs))
assert len(unexpected) == 0
# 1.2 T2I pipeline = AnimationPipeline(
if model_config.path != "": vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet,
if model_config.path.endswith(".ckpt"): scheduler=DDIMScheduler(**OmegaConf.to_container(inference_config.noise_scheduler_kwargs)),
state_dict = torch.load(model_config.path) ).to("cuda")
pipeline.unet.load_state_dict(state_dict)
elif model_config.path.endswith(".safetensors"): # 1. unet ckpt
state_dict = {} # 1.1 motion module
with safe_open(model_config.path, framework="pt", device="cpu") as f: motion_module_state_dict = torch.load(motion_module, map_location="cpu")
for key in f.keys(): if "global_step" in motion_module_state_dict: func_args.update({"global_step": motion_module_state_dict["global_step"]})
state_dict[key] = f.get_tensor(key) missing, unexpected = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False)
assert len(unexpected) == 0
is_lora = all("lora" in k for k in state_dict.keys()) # 1.2 T2I
if not is_lora: if model_config.path != "":
base_state_dict = state_dict if model_config.path.endswith(".ckpt"):
else: state_dict = torch.load(model_config.path)
base_state_dict = {} pipeline.unet.load_state_dict(state_dict)
with safe_open(model_config.base, framework="pt", device="cpu") as f:
elif model_config.path.endswith(".safetensors"):
state_dict = {}
with safe_open(model_config.path, framework="pt", device="cpu") as f:
for key in f.keys(): for key in f.keys():
base_state_dict[key] = f.get_tensor(key) state_dict[key] = f.get_tensor(key)
# vae is_lora = all("lora" in k for k in state_dict.keys())
converted_vae_checkpoint = convert_ldm_vae_checkpoint(base_state_dict, pipeline.vae.config) if not is_lora:
pipeline.vae.load_state_dict(converted_vae_checkpoint) base_state_dict = state_dict
# unet else:
converted_unet_checkpoint = convert_ldm_unet_checkpoint(base_state_dict, pipeline.unet.config) base_state_dict = {}
pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False) with safe_open(model_config.base, framework="pt", device="cpu") as f:
# text_model for key in f.keys():
pipeline.text_encoder = convert_ldm_clip_checkpoint(base_state_dict) base_state_dict[key] = f.get_tensor(key)
# import pdb # vae
# pdb.set_trace() converted_vae_checkpoint = convert_ldm_vae_checkpoint(base_state_dict, pipeline.vae.config)
if is_lora: pipeline.vae.load_state_dict(converted_vae_checkpoint)
pipeline = convert_lora(pipeline, state_dict, alpha=model_config.lora_alpha) # unet
converted_unet_checkpoint = convert_ldm_unet_checkpoint(base_state_dict, pipeline.unet.config)
pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False)
# text_model
pipeline.text_encoder = convert_ldm_clip_checkpoint(base_state_dict)
pipeline.to("cuda") # import pdb
### <<< create validation pipeline <<< ### # pdb.set_trace()
if is_lora:
pipeline = convert_lora(pipeline, state_dict, alpha=model_config.lora_alpha)
prompts = model_config.prompt pipeline.to("cuda")
n_prompts = list(model_config.n_prompt) * len(prompts) if len(model_config.n_prompt) == 1 else model_config.n_prompt ### <<< create validation pipeline <<< ###
random_seeds = model_config.pop("seed", [-1]) prompts = model_config.prompt
random_seeds = [random_seeds] if isinstance(random_seeds, int) else list(random_seeds) n_prompts = list(model_config.n_prompt) * len(prompts) if len(model_config.n_prompt) == 1 else model_config.n_prompt
random_seeds = random_seeds * len(prompts) if len(random_seeds) == 1 else random_seeds
config[config_key].random_seed = [] random_seeds = model_config.get("seed", [-1])
for prompt_idx, (prompt, n_prompt, random_seed) in enumerate(zip(prompts, n_prompts, random_seeds)): random_seeds = [random_seeds] if isinstance(random_seeds, int) else list(random_seeds)
random_seeds = random_seeds * len(prompts) if len(random_seeds) == 1 else random_seeds
# manually set random seed for reproduction config[config_key].random_seed = []
if random_seed != -1: torch.manual_seed(random_seed) for prompt_idx, (prompt, n_prompt, random_seed) in enumerate(zip(prompts, n_prompts, random_seeds)):
else: torch.seed()
config[config_key].random_seed.append(torch.initial_seed())
print(f"current seed: {torch.initial_seed()}") # manually set random seed for reproduction
print(f"sampling {prompt} ...") if random_seed != -1: torch.manual_seed(random_seed)
sample = pipeline( else: torch.seed()
prompt, config[config_key].random_seed.append(torch.initial_seed())
negative_prompt = n_prompt,
num_inference_steps = model_config.steps,
guidance_scale = model_config.guidance_scale,
width = args.W,
height = args.H,
video_length = args.L,
).videos
samples.append(sample)
prompt = "-".join((prompt.replace("/", "").split(" ")[:10])) print(f"current seed: {torch.initial_seed()}")
save_videos_grid(sample, f"{savedir}/sample/{model_idx}-{prompt_idx}-{prompt}.gif") print(f"sampling {prompt} ...")
print(f"save to {savedir}/sample/{prompt}.gif") sample = pipeline(
prompt,
negative_prompt = n_prompt,
num_inference_steps = model_config.steps,
guidance_scale = model_config.guidance_scale,
width = args.W,
height = args.H,
video_length = args.L,
).videos
samples.append(sample)
prompt = "-".join((prompt.replace("/", "").split(" ")[:10]))
save_videos_grid(sample, f"{savedir}/sample/{sample_idx}-{prompt}.gif")
print(f"save to {savedir}/sample/{prompt}.gif")
sample_idx += 1
samples = torch.concat(samples) samples = torch.concat(samples)
save_videos_grid(samples, f"{savedir}/sample.gif", n_rows=4) save_videos_grid(samples, f"{savedir}/sample.gif", n_rows=4)