import yaml import tempfile import gradio as gr import os import torch import imageio import argparse from types import MethodType import safetensors.torch as sf import torch.nn.functional as F from omegaconf import OmegaConf from transformers import CLIPTextModel, CLIPTokenizer from diffusers import MotionAdapter, EulerAncestralDiscreteScheduler, AutoencoderKL from diffusers import AutoencoderKL, UNet2DConditionModel, DPMSolverMultistepScheduler from diffusers.models.attention_processor import AttnProcessor2_0 from torch.hub import download_url_to_file from src.ic_light import BGSource from src.animatediff_pipe import AnimateDiffVideoToVideoPipeline from src.ic_light_pipe import StableDiffusionImg2ImgPipeline from utils.tools import read_video from huggingface_hub import snapshot_download, hf_hub_download hf_hub_download( repo_id='lllyasviel/ic-light', filename='iclight_sd15_fc.safetensors', local_dir='./models' ) snapshot_download( repo_id="stablediffusionapi/realistic-vision-v51", local_dir="./models/stablediffusionapi/realistic-vision-v51" ) snapshot_download( repo_id="guoyww/animatediff-motion-adapter-v1-5-3", local_dir="./models/guoyww/animatediff-motion-adapter-v1-5-3" ) def main(args): config = OmegaConf.load(args.config) device = torch.device('cuda') adopted_dtype = torch.float16 set_all_seed(42) ## vdm model adapter = MotionAdapter.from_pretrained(args.motion_adapter_model) ## pipeline pipe = AnimateDiffVideoToVideoPipeline.from_pretrained(args.sd_model, motion_adapter=adapter) eul_scheduler = EulerAncestralDiscreteScheduler.from_pretrained( args.sd_model, subfolder="scheduler", beta_schedule="linear", ) pipe.scheduler = eul_scheduler pipe.enable_vae_slicing() pipe = pipe.to(device=device, dtype=adopted_dtype) pipe.vae.requires_grad_(False) pipe.unet.requires_grad_(False) ## ic-light model tokenizer = CLIPTokenizer.from_pretrained(args.sd_model, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(args.sd_model, subfolder="text_encoder") vae = AutoencoderKL.from_pretrained(args.sd_model, subfolder="vae") unet = UNet2DConditionModel.from_pretrained(args.sd_model, subfolder="unet") with torch.no_grad(): new_conv_in = torch.nn.Conv2d(8, unet.conv_in.out_channels, unet.conv_in.kernel_size, unet.conv_in.stride, unet.conv_in.padding) new_conv_in.weight.zero_() #torch.Size([320, 8, 3, 3]) new_conv_in.weight[:, :4, :, :].copy_(unet.conv_in.weight) new_conv_in.bias = unet.conv_in.bias unet.conv_in = new_conv_in unet_original_forward = unet.forward def hooked_unet_forward(sample, timestep, encoder_hidden_states, **kwargs): c_concat = kwargs['cross_attention_kwargs']['concat_conds'].to(sample) c_concat = torch.cat([c_concat] * (sample.shape[0] // c_concat.shape[0]), dim=0) new_sample = torch.cat([sample, c_concat], dim=1) kwargs['cross_attention_kwargs'] = {} return unet_original_forward(new_sample, timestep, encoder_hidden_states, **kwargs) unet.forward = hooked_unet_forward ## ic-light model loader if not os.path.exists(args.ic_light_model): download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors', dst=args.ic_light_model) sd_offset = sf.load_file(args.ic_light_model) sd_origin = unet.state_dict() sd_merged = {k: sd_origin[k] + sd_offset[k] for k in sd_origin.keys()} unet.load_state_dict(sd_merged, strict=True) del sd_offset, sd_origin, sd_merged text_encoder = text_encoder.to(device=device, dtype=adopted_dtype) vae = vae.to(device=device, dtype=adopted_dtype) unet = unet.to(device=device, dtype=adopted_dtype) unet.set_attn_processor(AttnProcessor2_0()) vae.set_attn_processor(AttnProcessor2_0()) # Consistent light attention @torch.inference_mode() def custom_forward_CLA(self, hidden_states, gamma=config.get("gamma", 0.5), encoder_hidden_states=None, attention_mask=None, cross_attention_kwargs=None ): batch_size, sequence_length, channel = hidden_states.shape residual = hidden_states input_ndim = hidden_states.ndim if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) if attention_mask is not None: if attention_mask.shape[-1] != query.shape[1]: target_length = query.shape[1] attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) attention_mask = attention_mask.repeat_interleave(self.heads, dim=0) if self.group_norm is not None: hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) if encoder_hidden_states is None: encoder_hidden_states = hidden_states query = self.to_q(hidden_states) key = self.to_k(encoder_hidden_states) value = self.to_v(encoder_hidden_states) inner_dim = key.shape[-1] head_dim = inner_dim // self.heads query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False) shape = query.shape # addition key and value mean_key = key.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True) mean_value = value.reshape(2,-1,shape[1],shape[2],shape[3]).mean(dim=1,keepdim=True) mean_key = mean_key.expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3]) mean_value = mean_value.expand(-1,shape[0]//2,-1,-1,-1).reshape(shape[0],shape[1],shape[2],shape[3]) add_hidden_state = F.scaled_dot_product_attention(query, mean_key, mean_value, attn_mask=None, dropout_p=0.0, is_causal=False) # mix hidden_states = (1-gamma)*hidden_states + gamma*add_hidden_state hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) hidden_states = hidden_states.to(query.dtype) hidden_states = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) if input_ndim == 4: hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) if self.residual_connection: hidden_states = hidden_states + residual hidden_states = hidden_states / self.rescale_output_factor return hidden_states ### attention @torch.inference_mode() def prep_unet_self_attention(unet): for name, module in unet.named_modules(): module_name = type(module).__name__ name_split_list = name.split(".") cond_1 = name_split_list[0] in "up_blocks" cond_2 = name_split_list[-1] in ('attn1') if "Attention" in module_name and cond_1 and cond_2: cond_3 = name_split_list[1] if cond_3 not in "3": module.forward = MethodType(custom_forward_CLA, module) return unet ## consistency light attention unet = prep_unet_self_attention(unet) ## ic-light-scheduler ic_light_scheduler = DPMSolverMultistepScheduler( num_train_timesteps=1000, beta_start=0.00085, beta_end=0.012, algorithm_type="sde-dpmsolver++", use_karras_sigmas=True, steps_offset=1 ) ic_light_pipe = StableDiffusionImg2ImgPipeline( vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, unet=unet, scheduler=ic_light_scheduler, safety_checker=None, requires_safety_checker=False, feature_extractor=None, image_encoder=None ) ic_light_pipe = ic_light_pipe.to(device) ############################# params ###################################### strength = config.get("strength", 0.5) num_step = config.get("num_step", 25) text_guide_scale = config.get("text_guide_scale", 2) seed = config.get("seed") image_width = config.get("width", 512) image_height = config.get("height", 512) n_prompt = config.get("n_prompt", "") relight_prompt = config.get("relight_prompt", "") video_path = config.get("video_path", "") bg_source = BGSource[config.get("bg_source")] save_path = config.get("save_path") ############################## infer ##################################### generator = torch.manual_seed(seed) video_name = os.path.basename(video_path) video_list, video_name = read_video(video_path, image_width, image_height) print("################## begin ##################") with torch.no_grad(): num_inference_steps = int(round(num_step / strength)) output = pipe( ic_light_pipe=ic_light_pipe, relight_prompt=relight_prompt, bg_source=bg_source, video=video_list, prompt=relight_prompt, strength=strength, negative_prompt=n_prompt, guidance_scale=text_guide_scale, num_inference_steps=num_inference_steps, height=image_height, width=image_width, generator=generator, ) frames = output.frames[0] results_path = f"{save_path}/relight_{video_name}" imageio.mimwrite(results_path, frames, fps=8) print(f"relight with bg generation! prompt:{relight_prompt}, light:{bg_source.value}, save in {results_path}.") def infer(n_prompt, relight_prompt, video_path, bg_source, save_path, width, height, strength, gamma, num_step, text_guide_scale, seed): config_data = { "n_prompt": n_prompt, "relight_prompt": relight_prompt, "video_path": video_path, "bg_source": bg_source, "save_path": save_path, "width": width, "height": height, "strength": strength, "gamma": gamma, "num_step": num_step, "text_guide_scale": text_guide_scale, "seed": seed } temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".yaml") with open(temp_file.name, 'w') as file: yaml.dump(config_data, file, default_flow_style=False) config_path = temp_file.name class Args: def __init__(self): self.sd_model = "./models/stablediffusionapi/realistic-vision-v51" self.motion_adapter_model = "./models/guoyww/animatediff-motion-adapter-v1-5-3" self.ic_light_model = "./models/iclight_sd15_fc.safetensors" self.config = config_path args = Args() main(args) video_name = os.path.basename(video_path) results_path = f"{save_path}/relight_{video_name}" os.remove(config_path) return results_path with gr.Blocks() as demo: with gr.Row(): n_prompt = gr.Textbox(label="Negative Prompt") relight_prompt = gr.Textbox(label="Relight Prompt") with gr.Row(): video_path = gr.Textbox(label="Video Path") bg_source = gr.Dropdown(["NONE", "LEFT", "RIGHT", "BOTTOM", "TOP"], label="Background Source") with gr.Row(): save_path = gr.Textbox(label="Save Path") width = gr.Number(label="Width", value=512) height = gr.Number(label="Height", value=512) with gr.Row(): strength = gr.Slider(minimum=0.0, maximum=1.0, label="Strength", value=0.5) gamma = gr.Slider(minimum=0.0, maximum=1.0, label="Gamma", value=0.5) with gr.Row(): num_step = gr.Number(label="Number of Steps", value=25) text_guide_scale = gr.Number(label="Text Guide Scale", value=2) seed = gr.Number(label="Seed", value=2060) output = gr.Textbox(label="Results Path") submit = gr.Button("Run") submit.click(infer, inputs=[n_prompt, relight_prompt, video_path, bg_source, save_path, width, height, strength, gamma, num_step, text_guide_scale, seed], outputs=output) demo.launch()