init

Light_A_Video_node.py

@@ -0,0 +1,158 @@
+# !/usr/bin/env python
+# -*- coding: UTF-8 -*-
+import os
+import torch
+import gc
+import numpy as np
+from diffusers import StableDiffusionPipeline
+from .lav_relight import load_ic_light_model,infer_relight
+from .node_utils import load_images,tensor2pil_list
+import folder_paths
+from .src.ic_light import BGSource
+from .src.tools import  set_all_seed
+
+MAX_SEED = np.iinfo(np.int32).max
+current_node_path = os.path.dirname(os.path.abspath(__file__))
+device = torch.device(
+    "cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+# add checkpoints dir
+Light_A_Video_weigths_path = os.path.join(folder_paths.models_dir, "Light_A_Video")
+if not os.path.exists(Light_A_Video_weigths_path):
+    os.makedirs(Light_A_Video_weigths_path)
+folder_paths.add_model_folder_path("Light_A_Video", Light_A_Video_weigths_path)
+
+
+class Light_A_Video_Loader:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": (folder_paths.get_filename_list("checkpoints"),),
+                "motion_adapter_model": (["none"] + folder_paths.get_filename_list("controlnet"),),
+                "ic_light_model": (["none"] + folder_paths.get_filename_list("controlnet"),),
+                "mode":(["relight","inpaint"],),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL_Light_A_Video",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "loader_main"
+    CATEGORY = "Light_A_Video"
+
+    def loader_main(self, model,motion_adapter_model, ic_light_model,mode):
+
+        adopted_dtype = torch.float16
+
+        original_config_file=os.path.join(folder_paths.models_dir,"configs","v1-inference.yaml")
+        sd_repo = os.path.join(current_node_path, "sd_repo")
+        motion_repo=os.path.join(current_node_path,"animate_repo")
+        if model!="none":
+            ckpt_path=folder_paths.get_full_path("checkpoints",model)
+        else:
+            raise "no checkpoint"
+        try:
+            pipeline = StableDiffusionPipeline.from_single_file(
+            ckpt_path,config=sd_repo, original_config=original_config_file)
+        except:
+            pipeline = StableDiffusionPipeline.from_single_file(
+            ckpt_path, config=sd_repo,original_config_file=original_config_file)
+
+        # load model
+        print("***********Load model ***********")
+        motion_adapter_model=folder_paths.get_full_path("controlnet",motion_adapter_model)
+        ic_light_model=folder_paths.get_full_path("controlnet",ic_light_model)
+    
+        pipe,ic_light_pipe=load_ic_light_model(pipeline,ic_light_model,ckpt_path,sd_repo,motion_repo,motion_adapter_model,device,adopted_dtype,mode)
+        print("***********Load model done ***********")
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        return ({"model":pipe,"ic_light_pipe":ic_light_pipe,"mode":mode,"adopted_dtype":adopted_dtype},)
+
+
+class Light_A_Video_Sampler:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL_Light_A_Video",),
+                "images": ("IMAGE",),
+                "relight_prompt": ("STRING", {"default": "a car driving on the street, neon light", "multiline": True}),
+                "inpaint_prompt": ("STRING", {"default": "a car driving on the beach, sunset over sea", "multiline": True}),
+                "n_prompt": ("STRING", {"default": "bad quality, worse quality", "multiline": True}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
+                "num_step": ("INT", {"default": 25, "min": 1, "max": 1024, "step": 1, "display": "number"}),
+                "strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
+                "text_guide_scale": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1, "display": "number"}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "number"}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "number"}),
+                "bg_target": (["LEFT", "NONE", "RIGHT", "TOP", "BOTTOM"],),
+                "mask_repo": ("STRING", {"default": "ZhengPeng7/BiRefNet"},),},
+            "optional": {"mask_img": ("IMAGE",),
+                         "fps": ("FLOAT", {"default": 8.0, "min": 8.0, "max": 100.0, "step": 0.1}),
+                         },
+            }
+
+    RETURN_TYPES = ("IMAGE", )
+    RETURN_NAMES = ("image",)
+    FUNCTION = "sampler_main"
+    CATEGORY = "Light_A_Video"
+
+    def sampler_main(self, model,images,relight_prompt,inpaint_prompt, n_prompt,seed, num_step, strength,text_guide_scale,width, height,bg_target,mask_repo,**kwargs):
+        set_all_seed(42)
+
+        local_sam=os.path.join(Light_A_Video_weigths_path,"sam2_b.pt")
+        if os.path.exists(local_sam):
+            local_sam="sam2_b.pt"
+        ref_image_list=tensor2pil_list(images,width,height)
+
+        if bg_target=="NONE":
+            bg_source = BGSource.NONE
+        elif bg_target=="LEFT":
+            bg_source = BGSource.LEFT
+        elif bg_target=="RIGHT":
+            bg_source = BGSource.RIGHT
+        elif bg_target=="TOP":
+            bg_source = BGSource.TOP
+        else:
+            bg_source = BGSource.BOTTOM
+
+        mask_img=kwargs.get("mask_img",None)
+        fps=kwargs.get("fps",8.0)
+        ic_light_pipe=model.get("ic_light_pipe")
+        pipe=model.get("model")
+        mode=model.get("mode")
+        adopted_dtype=model.get("adopted_dtype")
+
+        if isinstance(mask_img,torch.Tensor): 
+            mask_list=tensor2pil_list(mask_img,width,height)
+            if len(mask_list)==1:
+                mask_list=mask_list*len(ref_image_list)
+                print("not enough mask, repeat mask to match the number of images")
+        else:
+            mask_list=None
+        
+        print("***********Start infer  ***********")
+        iamge = infer_relight(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,width,height,n_prompt,relight_prompt,inpaint_prompt,ref_image_list,bg_source,mode,mask_list,device,adopted_dtype,mask_repo,fps,local_sam)
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        return (load_images(iamge),)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Light_A_Video_Loader": Light_A_Video_Loader,
+    "Light_A_Video_Sampler": Light_A_Video_Sampler,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Light_A_Video_Loader": "Light_A_Video_Loader",
+    "Light_A_Video_Sampler": "Light_A_Video_Sampler",
+}

__init__.py

@@ -0,0 +1,4 @@
+
+from .Light_A_Video_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

animate_repo/config.json

@@ -0,0 +1,16 @@
+{
+  "_class_name": "MotionAdapter",
+  "_diffusers_version": "0.25.0.dev0",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "motion_layers_per_block": 2,
+  "motion_max_seq_length": 32,
+  "motion_mid_block_layers_per_block": 1,
+  "motion_norm_num_groups": 32,
+  "motion_num_attention_heads": 8,
+  "use_motion_mid_block": false
+}

configs/relight/bear.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a bear walking on the rock, nature lighting, key light"
+video_path: "input_animatediff/bear.mp4"
+bg_source: "TOP"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/boat.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a boat floating on the sea, sunset"
+video_path: "input_animatediff/boat.mp4"
+bg_source: "TOP"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/car.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a car driving on the street, neon light"
+video_path: "input_animatediff/car.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 2060

configs/relight/cat.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a cat, red and blue neon light"
+video_path: "input_animatediff/cat.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/cow.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a cow drinking water in the river, sunset"
+video_path: "input_animatediff/cow.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/flowers.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality, unclear, blurry"
+relight_prompt: "A basket of flowers, sunshine, hard light"
+video_path: "input_animatediff/flowers.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/fox.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a fox, sunlight filtering through trees, dappled light"
+video_path: "input_animatediff/fox.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/girl.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a girl, magic lit, sci-fi RGB glowing, key lighting"
+video_path: "input_animatediff/girl.mp4"
+bg_source: "BOTTOM"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/girl2.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "an anime girl, neon light"
+video_path: "input_animatediff/girl2.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/juice.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "Pour juice into a glass, magic golden lit"
+video_path: "input_animatediff/juice.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/man2.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "handsome man with glasses, shadow from window, sunshine"
+video_path: "input_animatediff/man2.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/man4.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "handsome man with glasses, sunlight through the blinds"
+video_path: "input_animatediff/man4.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/plane.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a plane on the runway, bottom neon light"
+video_path: "input_animatediff/plane.mp4"
+bg_source: "BOTTOM"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/toy.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a maneki-neko toy, cozy bedroom illumination"
+video_path: "input_animatediff/toy.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight/woman.yaml

@@ -0,0 +1,13 @@
+n_prompt: "bad quality, worse quality"
+relight_prompt: "a woman with curly hair, natural lighting, warm atmosphere"
+video_path: "input_animatediff/woman.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 25
+text_guide_scale: 2
+seed: 42

configs/relight_inpaint/bloom.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a red flower blooming in the river"
+relight_prompt: "a red flower blooming in the river, nature lighting"
+
+video_path: "input_animatediff/bloom.mp4"
+bg_source: "TOP"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.4
+gamma: 0.5
+num_step: 50
+text_guide_scale: 4
+seed: 8776

configs/relight_inpaint/camera.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "A tiny camera on a tray, cyberpunk"
+relight_prompt: "A tiny camera on a tray, cyberpunk, neon light"
+
+video_path: "input_animatediff/camera.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.4
+gamma: 0.5
+num_step: 50
+text_guide_scale: 3
+seed: 1333

configs/relight_inpaint/car.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a car driving on the street"
+relight_prompt: "a car driving on the street, neon light"
+
+video_path: "input_animatediff/car.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 2
+seed: 6561

configs/relight_inpaint/car_2.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a car driving on the beach, sunset over sea"
+relight_prompt: "a car driving on the beach, sunset over sea, left light, shadow"
+
+video_path: "input_animatediff/car.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 2
+seed: 2409

configs/relight_inpaint/cat2.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "A cat walking on a runway, red and blue neon lights on both sides"
+relight_prompt: "A cat walking on a runway, red and blue neon lights on both sides, key light"
+
+video_path: "input_animatediff/cat2.mp4"
+bg_source: "LEFT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 5
+seed: 2949

configs/relight_inpaint/coin.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "A coin on the desk"
+relight_prompt: "A coin on the desk, natural lighting"
+
+video_path: "input_animatediff/coin.mp4"
+bg_source: "TOP"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.4
+gamma: 0.5
+num_step: 80
+text_guide_scale: 2
+seed: 4013

configs/relight_inpaint/dog2.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a dog in the room, sunshine from window"
+relight_prompt: "a dog in the room, sunshine from window"
+
+video_path: "input_animatediff/dog2.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.4
+gamma: 0.5
+num_step: 50
+text_guide_scale: 2
+seed: 4550

configs/relight_inpaint/man3.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "A man in the classroom"
+relight_prompt: "A man in the classroom, sunshine from the window"
+
+video_path: "input_animatediff/man3.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 3
+seed: 3931

configs/relight_inpaint/man3_2.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "A man in a bar, left yellow and right purple neon lights"
+relight_prompt: "A man in a bar, left yellow and right purple neon lights, hard light"
+
+video_path: "input_animatediff/man3.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 4
+seed: 9528

configs/relight_inpaint/water.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a glass of water, in the forest, magic golden lit"
+relight_prompt: "a glass of water, in the forest, magic golden lit, key light"
+
+video_path: "input_animatediff/water.mp4"
+bg_source: "TOP"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.4
+gamma: 0.5
+num_step: 50
+text_guide_scale: 4
+seed: 796

configs/relight_inpaint/wolf2.yaml

@@ -0,0 +1,15 @@
+n_prompt: "bad quality, worse quality"
+inpaint_prompt: "a wolf stands in an alley, detailed face, neon, Wong Kar-wai, warm"
+relight_prompt: "a wolf stands in an alley, detailed face, neon, Wong Kar-wai, warm, right light"
+
+video_path: "input_animatediff/wolf2.mp4"
+bg_source: "RIGHT"  ## NONE, LEFT, RIGHT, BOTTOM, TOP
+save_path: "output"
+
+width: 512
+height: 512
+strength: 0.5
+gamma: 0.5
+num_step: 50
+text_guide_scale: 5
+seed: 2172

lav_paint.py

@@ -0,0 +1,256 @@
+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 .src.ic_light import BGSource
+from .src.ic_light import Relighter
+from .src.animatediff_inpaint_pipe import AnimateDiffVideoToVideoPipeline
+from .src.ic_light_pipe import StableDiffusionImg2ImgPipeline
+from .src.tools import read_video, read_mask,set_all_seed, get_fg_video
+
+
+
+
+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)
+    #MotionAdapter.from_single_file(args.motion_adapter_model, torch_dtype=adopted_dtype)
+    ## 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
+    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", 50)
+    text_guide_scale = config.get("text_guide_scale", 4)
+    seed = config.get("seed")
+    image_width = config.get("width", 512)
+    image_height = config.get("height", 512)
+    n_prompt = config.get("n_prompt", "")
+    inpaint_prompt = config.get("inpaint_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)
+    mask_folder = os.path.join("masks_animatediff", video_name.split('.')[-2])
+    mask_list = read_mask(mask_folder)
+
+    print("################## begin ##################")
+    ## get foreground video
+    fg_video_tensor = get_fg_video(video_list, mask_list, device, adopted_dtype) ## torch.Size([16, 3, 512, 512])
+
+    with torch.no_grad():
+        relighter = Relighter(
+            pipeline=ic_light_pipe,
+            relight_prompt=relight_prompt,
+            bg_source=bg_source, 
+            generator=generator,
+            )
+        vdm_init_latent = relighter(fg_video_tensor)
+
+        ## infer
+        num_inference_steps = num_step
+        output = pipe(
+            ic_light_pipe=ic_light_pipe,
+            relight_prompt=relight_prompt,
+            bg_source=bg_source,
+            mask=mask_list,
+            vdm_init_latent=vdm_init_latent,
+            video=video_list,
+            prompt=inpaint_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}/inpaint_{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}.")
+    
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument("--sd_model", type=str, default="stablediffusionapi/realistic-vision-v51")
+    parser.add_argument("--motion_adapter_model", type=str, default="guoyww/animatediff-motion-adapter-v1-5-3")
+    parser.add_argument("--ic_light_model", type=str, default="./models/iclight_sd15_fc.safetensors")
+    
+    parser.add_argument("--config", type=str, default="configs/relight_inpaint/car.yaml", help="the config file for each sample.")
+    
+    args = parser.parse_args()
+    main(args)

lav_relight.py

@@ -0,0 +1,252 @@
+import os
+import torch
+from types import MethodType
+import safetensors.torch as sf
+import torch.nn.functional as F
+from transformers import CLIPTokenizer
+from diffusers import MotionAdapter, EulerAncestralDiscreteScheduler
+from diffusers import  UNet2DConditionModel, DPMSolverMultistepScheduler
+from diffusers.models.attention_processor import AttnProcessor2_0
+from torch.hub import download_url_to_file
+from .src.tools import get_fg_video
+from .src.animatediff_pipe import AnimateDiffVideoToVideoPipeline
+from .src.animatediff_inpaint_pipe import AnimateDiffVideoToVideoPipeline as AnimateDiffVideoToVideoPipeline_inpaint
+from .src.ic_light_pipe import StableDiffusionImg2ImgPipeline
+from .src.ic_light import Relighter
+from .node_utils import image2masks
+import logging
+logging.basicConfig(level=logging.INFO)
+
+
+def load_ic_light_model(pipeline,ic_light_model,ckpt_path,sd_repo,motion_repo,motion_adapter_model,device,adopted_dtype,mode):
+     # vdm model
+    adapter = MotionAdapter.from_single_file(motion_adapter_model,config=motion_repo)
+    Unet=UNet2DConditionModel.from_single_file(ckpt_path, config=os.path.join(sd_repo, "unet"))
+    # animate pipeline
+    Vae=pipeline.vae
+    tokenizer = CLIPTokenizer.from_pretrained(sd_repo, subfolder="tokenizer")
+    Text_encoder =pipeline.text_encoder
+    del pipeline
+
+    if mode=='relight':
+        pipe=AnimateDiffVideoToVideoPipeline.from_pretrained(sd_repo,unet=Unet,vae=Vae,text_encoder=Text_encoder,tokenizer=tokenizer,motion_adapter=adapter)
+    else:
+        pipe=AnimateDiffVideoToVideoPipeline_inpaint.from_pretrained(sd_repo,unet=Unet,vae=Vae,text_encoder=Text_encoder,tokenizer=tokenizer,motion_adapter=adapter)
+    eul_scheduler = EulerAncestralDiscreteScheduler.from_pretrained(
+        sd_repo,
+        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)
+
+    unet=UNet2DConditionModel.from_single_file(ckpt_path,config=os.path.join(sd_repo, "unet"))
+    tokenizer=pipe.tokenizer
+    text_encoder=pipe.text_encoder
+    vae=Vae
+
+
+    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
+
+    if not os.path.exists(ic_light_model):
+        download_url_to_file(url='https://huggingface.co/lllyasviel/ic-light/resolve/main/iclight_sd15_fc.safetensors', 
+                             dst=ic_light_model) #TO DO: change to ic-light model
+    
+    sd_offset = sf.load_file(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())
+    
+    @torch.inference_mode()
+    def custom_forward_CLA(self, 
+                        hidden_states, 
+                        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)
+    return pipe,ic_light_pipe
+
+
+def infer_relight(ic_light_pipe,pipe,strength,num_step,text_guide_scale,seed,image_width,image_height,n_prompt,relight_prompt,inpaint_prompt,video_list,bg_source,mode,mask_list,device,adopted_dtype,repo,fps,local_sam):
+
+    generator = torch.manual_seed(seed)
+    if mode == "inpaint":
+        if mask_list is None:
+            if repo:
+                print("No mask_list provided, generating mask_list from repo...")
+                mask_list = image2masks(repo,video_list)
+            else:
+                from .sam2 import get_mask
+                mask_list=get_mask(video_list,fps,local_sam)
+                print("No mask_list provided, use sam2 generating mask_list from video...")
+        fg_video_tensor = get_fg_video(video_list, mask_list, device, adopted_dtype)
+        with torch.no_grad():
+            relighter = Relighter(
+                pipeline=ic_light_pipe,
+                relight_prompt=relight_prompt,
+                bg_source=bg_source, 
+                generator=generator,
+                )
+            vdm_init_latent = relighter(fg_video_tensor)
+
+            ## infer
+            num_inference_steps = num_step
+            output = pipe(
+                ic_light_pipe=ic_light_pipe,
+                relight_prompt=relight_prompt,
+                bg_source=bg_source,
+                mask=mask_list,
+                vdm_init_latent=vdm_init_latent,
+                video=video_list,
+                prompt=inpaint_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,
+            )
+    else:
+        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]
+    return frames
+

node_utils.py

@@ -0,0 +1,241 @@
+# !/usr/bin/env python
+# -*- coding: UTF-8 -*-
+import os
+import torch
+from PIL import Image
+import numpy as np
+import cv2
+import gc
+import time
+from transformers import AutoModelForImageSegmentation
+from comfy.utils import common_upscale,ProgressBar
+from huggingface_hub import hf_hub_download
+import torchvision.transforms as transforms
+cur_path = os.path.dirname(os.path.abspath(__file__))
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+def cv2pil(cv_image):
+    """
+    将OpenCV图像转换为PIL图像
+    :param cv_image: OpenCV图像
+    :return: PIL图像
+    """
+    # 将图像从BGR转换为RGB
+    rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
+    # 使用PIL的Image.fromarray方法将NumPy数组转换为PIL图像
+    pil_image = Image.fromarray(rgb_image)
+    return pil_image
+
+def image2masks(repo,video_image):
+    start_time = time.time()
+    model = AutoModelForImageSegmentation.from_pretrained(repo, trust_remote_code=True)
+    torch.set_float32_matmul_precision(['high', 'highest'][0])
+    model.to('cuda')
+    model.eval()
+    # Data settings
+    image_size = (1024, 1024)
+    transform_image = transforms.Compose([
+    transforms.Resize(image_size),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    masks=[]
+    for img in video_image:
+        input_images = transform_image(img).unsqueeze(0).to('cuda')
+        # Prediction
+        with torch.no_grad():
+            preds = model(input_images)[-1].sigmoid().cpu()
+        pred = preds[0].squeeze()
+        pred_pil = transforms.ToPILImage()(pred)
+        mask = pred_pil.resize(img.size)
+        #img.putalpha(mask)
+        masks.append(mask.convert('RGB'))
+    end_time = time.time()  
+    load_time = end_time - start_time  
+    print(f"image2masks infer time: {load_time:.4f} s")
+    model.to('cpu')
+    gc.collect()
+    torch.cuda.empty_cache()
+    return masks
+
+def tensor_to_pil(tensor):
+    image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
+    image = Image.fromarray(image_np, mode='RGB')
+    return image
+
+def tensor2pil_list(image,width,height):
+    B,_,_,_=image.size()
+    if  B==1:
+        ref_image_list=[tensor2pil_upscale(image,width,height)]
+    else:
+        img_list = list(torch.chunk(image, chunks=B))
+        ref_image_list = [tensor2pil_upscale(img,width,height) for img in img_list]
+    return ref_image_list
+
+
+def tensor_upscale(img_tensor, width, height):
+    samples = img_tensor.movedim(-1, 1)
+    img = common_upscale(samples, width, height, "nearest-exact", "center")
+    samples = img.movedim(1, -1)
+    return samples
+
+def tensor2pil_upscale(img_tensor, width, height):
+    samples = img_tensor.movedim(-1, 1)
+    img = common_upscale(samples, width, height, "nearest-exact", "center")
+    samples = img.movedim(1, -1)
+    img_pil = tensor_to_pil(samples)
+    return img_pil
+
+
+def tensor2cv(tensor_image,RGB2BGR=True):
+    if len(tensor_image.shape)==4:#bhwc to hwc
+        tensor_image=tensor_image.squeeze(0)
+    if tensor_image.is_cuda:
+        tensor_image = tensor_image.cpu().detach()
+    tensor_image=tensor_image.numpy()
+    #反归一化
+    maxValue=tensor_image.max()
+    tensor_image=tensor_image*255/maxValue
+    img_cv2=np.uint8(tensor_image)#32 to uint8
+    if RGB2BGR:
+        img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR)
+    return img_cv2
+
+def cvargb2tensor(img):
+    assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
+    img = torch.from_numpy(img.transpose((2, 0, 1)))
+    return img.float().div(255).unsqueeze(0)  # 255也可以改为256
+
+def cv2tensor(img):
+    assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = torch.from_numpy(img.transpose((2, 0, 1)))
+    return img.float().div(255).unsqueeze(0)  # 255也可以改为256
+
+def images_generator(img_list: list,):
+    #get img size
+    sizes = {}
+    for image_ in img_list:
+        if isinstance(image_,Image.Image):
+            count = sizes.get(image_.size, 0)
+            sizes[image_.size] = count + 1
+        elif isinstance(image_,np.ndarray):
+            count = sizes.get(image_.shape[:2][::-1], 0)
+            sizes[image_.shape[:2][::-1]] = count + 1
+        else:
+            raise "unsupport image list,must be pil or cv2!!!"
+    size = max(sizes.items(), key=lambda x: x[1])[0]
+    yield size[0], size[1]
+    
+    # any to tensor
+    def load_image(img_in):
+        if isinstance(img_in, Image.Image):
+            img_in=img_in.convert("RGB")
+            i = np.array(img_in, dtype=np.float32)
+            i = torch.from_numpy(i).div_(255)
+            if i.shape[0] != size[1] or i.shape[1] != size[0]:
+                i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
+                i = common_upscale(i, size[0], size[1], "lanczos", "center")
+                i = i.squeeze(0).movedim(0, -1).numpy()
+            return i
+        elif isinstance(img_in,np.ndarray):
+            i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32)
+            i = torch.from_numpy(i).div_(255)
+            #print(i.shape)
+            return i
+        else:
+           raise "unsupport image list,must be pil,cv2 or tensor!!!"
+        
+    total_images = len(img_list)
+    processed_images = 0
+    pbar = ProgressBar(total_images)
+    images = map(load_image, img_list)
+    try:
+        prev_image = next(images)
+        while True:
+            next_image = next(images)
+            yield prev_image
+            processed_images += 1
+            pbar.update_absolute(processed_images, total_images)
+            prev_image = next_image
+    except StopIteration:
+        pass
+    if prev_image is not None:
+        yield prev_image
+
+def load_images(img_list: list,):
+    gen = images_generator(img_list)
+    (width, height) = next(gen)
+    images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
+    if len(images) == 0:
+        raise FileNotFoundError(f"No images could be loaded .")
+    return images
+
+def tensor2pil(tensor):
+    image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
+    image = Image.fromarray(image_np, mode='RGB')
+    return image
+
+def pil2narry(img):
+    narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
+    return narry
+
+def equalize_lists(list1, list2):
+    """
+    比较两个列表的长度，如果不一致，则将较短的列表复制以匹配较长列表的长度。
+    
+    参数:
+    list1 (list): 第一个列表
+    list2 (list): 第二个列表
+    
+    返回:
+    tuple: 包含两个长度相等的列表的元组
+    """
+    len1 = len(list1)
+    len2 = len(list2)
+    
+    if len1 == len2:
+        pass
+    elif len1 < len2:
+        print("list1 is shorter than list2, copying list1 to match list2's length.")
+        list1.extend(list1 * ((len2 // len1) + 1))  # 复制list1以匹配list2的长度
+        list1 = list1[:len2]  # 确保长度一致
+    else:
+        print("list2 is shorter than list1, copying list2 to match list1's length.")
+        list2.extend(list2 * ((len1 // len2) + 1))  # 复制list2以匹配list1的长度
+        list2 = list2[:len1]  # 确保长度一致
+    
+    return list1, list2
+
+def file_exists(directory, filename):
+    # 构建文件的完整路径
+    file_path = os.path.join(directory, filename)
+    # 检查文件是否存在
+    return os.path.isfile(file_path)
+
+def download_weights(file_dir,repo_id,subfolder="",pt_name=""):
+    if subfolder:
+        file_path = os.path.join(file_dir,subfolder, pt_name)
+        sub_dir=os.path.join(file_dir,subfolder)
+        if not os.path.exists(sub_dir):
+            os.makedirs(sub_dir)
+        if not os.path.exists(file_path):
+            file_path = hf_hub_download(
+                repo_id=repo_id,
+                subfolder=subfolder,
+                filename=pt_name,
+                local_dir = file_dir,
+            )
+        return file_path
+    else:
+        file_path = os.path.join(file_dir, pt_name)
+        if not os.path.exists(file_dir):
+            os.makedirs(file_dir)
+        if not os.path.exists(file_path):
+            file_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=pt_name,
+                local_dir=file_dir,
+            )
+        return file_path

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "light_a_video"
+description = "Light-A-Video: Training-free Video Relighting via Progressive Light Fusion,you can use it in comfyUI"
+version = "1.0.0"
+license = {file = "LICENSE"}
+dependencies = ["#diffusers==0.32.1", "diffusers", "#transformers==4.47.1", "transformers", "opencv-python", "safetensors", "pillow", "einops", "peft", "imageio", "omegaconf", "#ultralytic", "tqdm==4.67.1", "protobuf==3.20", "torch", "torchvision", "moviepy"]
+
+[project.urls]
+Repository = "https://github.com/smthemex/ComfyUI_Light_A_Video"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "smthemx"
+DisplayName = "ComfyUI_Light_A_Video"
+Icon = ""

requirements.txt

@@ -0,0 +1,17 @@
+#diffusers==0.32.1
+diffusers
+#transformers==4.47.1
+transformers
+opencv-python
+safetensors
+pillow
+einops
+peft
+imageio
+omegaconf
+#ultralytic
+tqdm==4.67.1
+protobuf==3.20
+torch
+torchvision
+moviepy

sam2.py

@@ -0,0 +1,32 @@
+import os
+import torch
+#import argparse
+import numpy as np
+from PIL import Image
+import random
+import imageio
+import folder_paths
+
+
+def get_mask(video_list,fps,local_sam, x=255, y=255):
+    from ultralytics.models.sam import SAM2VideoPredictor #ultralytics>=8.3.0
+    overrides = dict(conf=0.25, task="segment", mode="predict", imgsz=1024, model=local_sam)
+    predictor = SAM2VideoPredictor(overrides=overrides)
+    file_prefix = ''.join(random.choice("0123456789") for _ in range(6))
+    frames = [np.array(img) for img in video_list]
+    video_file = os.path.join(folder_paths.get_input_directory(), f"audio_{file_prefix}_temp.mp4")
+    imageio.mimsave(video_file, frames, fps=fps, codec='libx264')
+    results = predictor(source=video_file,points=[x, y],labels=[1])
+    mask_list=[]
+    for i in range(len(results)):
+        mask = (results[i].masks.data).squeeze().to(torch.float16)
+        mask = (mask * 255).cpu().numpy().astype(np.uint8)
+        mask_image = Image.fromarray(mask).convert('RGB')
+        mask_list.append(mask_image)
+        # mask_dir = f'masks_animatediff/{video_name}'
+        # if not os.path.exists(mask_dir):  
+        #     os.makedirs(mask_dir)        
+        mask_image.save(folder_paths.get_input_directory() + f'/{str(i).zfill(3)}.png')
+    return mask_list
+
+

sd_repo/README.md

@@ -0,0 +1,64 @@
+---
+license: creativeml-openrail-m
+tags:
+  - modelslab.com
+  - stable-diffusion-api
+  - text-to-image
+  - ultra-realistic
+pinned: true
+---
+
+#  API Inference
+
+![generated from modelslab.com](https://cdn2.stablediffusionapi.com/generations/bf190b5a-fe19-437c-ba05-82f29cb1f7ad-0.png)
+## Get API Key
+
+Get API key from [ModelsLab API](http://modelslab.com), No Payment needed. 
+
+Replace Key in below code, change **model_id**  to "realistic-vision-v51"
+
+Coding in PHP/Node/Java etc? Have a look at docs for more code examples: [View docs](https://modelslab.com/docs)
+
+Try model for free: [Generate Images](https://modelslab.com/models/realistic-vision-v51)
+
+Model link: [View model](https://modelslab.com/models/realistic-vision-v51)
+
+View all models: [View Models](https://modelslab.com/models)
+
+    import requests  
+    import json  
+      
+    url =  "https://modelslab.com/api/v6/images/text2img"  
+      
+    payload = json.dumps({  
+    "key":  "your_api_key",  
+    "model_id":  "realistic-vision-v51",  
+    "prompt":  "ultra realistic close up portrait ((beautiful pale cyberpunk female with heavy black eyeliner)), blue eyes, shaved side haircut, hyper detail, cinematic lighting, magic neon, dark red city, Canon EOS R3, nikon, f/1.4, ISO 200, 1/160s, 8K, RAW, unedited, symmetrical balance, in-frame, 8K",  
+    "negative_prompt":  "painting, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, deformed, ugly, blurry, bad anatomy, bad proportions, extra limbs, cloned face, skinny, glitchy, double torso, extra arms, extra hands, mangled fingers, missing lips, ugly face, distorted face, extra legs, anime",  
+    "width":  "512",  
+    "height":  "512",  
+    "samples":  "1",  
+    "num_inference_steps":  "30",  
+    "safety_checker":  "no",  
+    "enhance_prompt":  "yes",  
+    "seed":  None,  
+    "guidance_scale":  7.5,  
+    "multi_lingual":  "no",  
+    "panorama":  "no",  
+    "self_attention":  "no",  
+    "upscale":  "no",  
+    "embeddings":  "embeddings_model_id",  
+    "lora":  "lora_model_id",  
+    "webhook":  None,  
+    "track_id":  None  
+    })  
+      
+    headers =  {  
+    'Content-Type':  'application/json'  
+    }  
+      
+    response = requests.request("POST", url, headers=headers, data=payload)  
+      
+    print(response.text)
+
+> Use this coupon code to get 25% off **DMGG0RBN** 

sd_repo/feature_extractor/preprocessor_config.json

@@ -0,0 +1,29 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "feature_extractor_type": "CLIPFeatureExtractor",
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPFeatureExtractor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  },
+  "use_square_size": false
+}

sd_repo/model_index.json

@@ -0,0 +1,37 @@
+{
+  "_class_name": "StableDiffusionPipeline",
+  "_diffusers_version": "0.25.0.dev0",
+  "feature_extractor": [
+    "transformers",
+    "CLIPFeatureExtractor"
+  ],
+  "image_encoder": [
+    null,
+    null
+  ],
+  "requires_safety_checker": true,
+  "safety_checker": [
+    "stable_diffusion",
+    "StableDiffusionSafetyChecker"
+  ],
+  "scheduler": [
+    "diffusers",
+    "PNDMScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

sd_repo/safety_checker/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "CompVis/stable-diffusion-safety-checker",
+  "architectures": [
+    "StableDiffusionSafetyChecker"
+  ],
+  "initializer_factor": 1.0,
+  "logit_scale_init_value": 2.6592,
+  "model_type": "clip",
+  "projection_dim": 768,
+  "text_config": {
+    "dropout": 0.0,
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "model_type": "clip_text_model",
+    "num_attention_heads": 12
+  },
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.2",
+  "vision_config": {
+    "dropout": 0.0,
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "model_type": "clip_vision_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "patch_size": 14
+  }
+}

sd_repo/scheduler/scheduler_config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "PNDMScheduler",
+  "_diffusers_version": "0.25.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "timestep_spacing": "leading",
+  "trained_betas": null
+}

sd_repo/text_encoder/config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.2",
+  "vocab_size": 49408
+}

sd_repo/tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

sd_repo/tokenizer/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

sd_repo/unet/config.json

@@ -0,0 +1,67 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.25.0.dev0",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

sd_repo/vae/config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.25.0.dev0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 512,
+  "scaling_factor": 0.18215,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

src/animatediff_eul.py

@@ -0,0 +1,94 @@
+import torch
+from typing import List, Optional, Tuple, Union
+
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    BaseOutput,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+class EulerAncestralDiscreteSchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+    pred_original_sample: Optional[torch.FloatTensor] = None
+
+
+def eul_step(
+    self,
+    model_output: torch.FloatTensor,
+    timestep: Union[float, torch.FloatTensor],
+    sample: torch.FloatTensor, 
+    fusion_latent,
+    pipe,
+    generator: Optional[torch.Generator] = None,
+    return_dict: bool = True,
+) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:
+
+    if (
+        isinstance(timestep, int)
+        or isinstance(timestep, torch.IntTensor)
+        or isinstance(timestep, torch.LongTensor)
+    ):
+        raise ValueError(
+            (
+                "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                " one of the `scheduler.timesteps` as a timestep."
+            ),
+        )
+
+    if self.step_index is None:
+        self._init_step_index(timestep)
+
+    sigma = self.sigmas[self.step_index]
+
+    # Upcast to avoid precision issues when computing prev_sample
+    sample = sample.to(torch.float32)
+
+    # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
+    if self.config.prediction_type == "epsilon": ## True, 计算x_0
+        pred_original_sample = sample - sigma * model_output
+    elif self.config.prediction_type == "v_prediction":
+        # * c_out + input * c_skip
+        pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
+    elif self.config.prediction_type == "sample":
+        raise NotImplementedError("prediction_type not implemented yet: sample")
+    else:
+        raise ValueError(
+            f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"
+        )
+        
+    ## fusion latent
+    pred_original_sample = fusion_latent
+    
+    sigma_from = self.sigmas[self.step_index]
+    sigma_to = self.sigmas[self.step_index + 1]
+    sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5
+    sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5
+
+    # 2. Convert to an ODE derivative
+    derivative = (sample - pred_original_sample) / sigma
+    dt = sigma_down - sigma
+
+    prev_sample = sample + derivative * dt
+
+    device = model_output.device
+    noise = randn_tensor(model_output.shape, dtype=model_output.dtype, device=device, generator=generator)
+
+    prev_sample = prev_sample + noise * sigma_up
+
+    # Cast sample back to model compatible dtype
+    prev_sample = prev_sample.to(model_output.dtype)
+
+    # upon completion increase step index by one
+    self._step_index += 1
+
+    if not return_dict:
+        return (prev_sample,)
+
+    return EulerAncestralDiscreteSchedulerOutput(
+        prev_sample=prev_sample, pred_original_sample=pred_original_sample
+    )

src/ic_light.py

@@ -0,0 +1,152 @@
+import torch
+import numpy as np
+from enum import Enum
+import math
+
+import torch.nn.functional as F
+from .tools import resize_and_center_crop, numpy2pytorch, pad, decode_latents, encode_video
+
+class BGSource(Enum):
+    NONE = "None"
+    LEFT = "Left Light"
+    RIGHT = "Right Light"
+    TOP = "Top Light"
+    BOTTOM = "Bottom Light"
+
+class Relighter:
+    def __init__(self, 
+                 pipeline, 
+                 relight_prompt="",
+                 num_frames=16,
+                 image_width=512,
+                 image_height=512, 
+                 num_samples=1, 
+                 steps=15, 
+                 cfg=2, 
+                 lowres_denoise=0.9, 
+                 bg_source=BGSource.RIGHT, 
+                 generator=None,
+                 ):
+        
+        self.pipeline = pipeline
+        self.image_width = image_width
+        self.image_height = image_height
+        self.num_samples = num_samples
+        self.steps = steps
+        self.cfg = cfg
+        self.lowres_denoise = lowres_denoise
+        self.bg_source = bg_source
+        self.generator = generator
+        self.device = pipeline.device
+        self.num_frames = num_frames
+        self.vae = self.pipeline.vae
+        
+        self.a_prompt = "best quality"
+        self.n_prompt = "lowres, bad anatomy, bad hands, cropped, worst quality"
+        positive_prompt = relight_prompt + ', ' + self.a_prompt
+        negative_prompt = self.n_prompt
+        tokenizer = self.pipeline.tokenizer
+        device = self.pipeline.device
+        vae = self.vae
+        
+        conds, unconds = self.encode_prompt_pair(tokenizer, device, positive_prompt, negative_prompt)
+        input_bg = self.create_background()
+        bg = resize_and_center_crop(input_bg, self.image_width, self.image_height)
+        bg_latent = numpy2pytorch([bg], device, vae.dtype)
+        bg_latent = vae.encode(bg_latent).latent_dist.mode() * vae.config.scaling_factor
+        
+        self.bg_latent = bg_latent.repeat(self.num_frames, 1, 1, 1) ## 固定光源
+        self.conds = conds.repeat(self.num_frames, 1, 1)
+        self.unconds = unconds.repeat(self.num_frames, 1, 1)
+        
+    def encode_prompt_inner(self, tokenizer, txt):
+        max_length = tokenizer.model_max_length
+        chunk_length = tokenizer.model_max_length - 2
+        id_start = tokenizer.bos_token_id
+        id_end = tokenizer.eos_token_id
+        id_pad = id_end
+
+        tokens = tokenizer(txt, truncation=False, add_special_tokens=False)["input_ids"]
+        chunks = [[id_start] + tokens[i: i + chunk_length] + [id_end] for i in range(0, len(tokens), chunk_length)]
+        chunks = [pad(ck, id_pad, max_length) for ck in chunks]
+
+        token_ids = torch.tensor(chunks).to(device=self.device, dtype=torch.int64)
+        conds = self.pipeline.text_encoder(token_ids).last_hidden_state
+        return conds
+
+    def encode_prompt_pair(self, tokenizer, device, positive_prompt, negative_prompt):
+        c = self.encode_prompt_inner(tokenizer, positive_prompt)
+        uc = self.encode_prompt_inner(tokenizer, negative_prompt)
+
+        c_len = float(len(c))
+        uc_len = float(len(uc))
+        max_count = max(c_len, uc_len)
+        c_repeat = int(math.ceil(max_count / c_len))
+        uc_repeat = int(math.ceil(max_count / uc_len))
+        max_chunk = max(len(c), len(uc))
+
+        c = torch.cat([c] * c_repeat, dim=0)[:max_chunk]
+        uc = torch.cat([uc] * uc_repeat, dim=0)[:max_chunk]
+
+        c = torch.cat([p[None, ...] for p in c], dim=1)
+        uc = torch.cat([p[None, ...] for p in uc], dim=1)
+
+        return c.to(device), uc.to(device)
+
+    def create_background(self):
+        
+        max_pix = 255
+        min_pix = 0
+        
+        print(f"max light pix:{max_pix}, min light pix:{min_pix}")
+        
+        if self.bg_source == BGSource.NONE:
+            return None
+        elif self.bg_source == BGSource.LEFT:
+            gradient = np.linspace(max_pix, min_pix, self.image_width)
+            image = np.tile(gradient, (self.image_height, 1))
+            return np.stack((image,) * 3, axis=-1).astype(np.uint8)
+        elif self.bg_source == BGSource.RIGHT:
+            gradient = np.linspace(min_pix, max_pix, self.image_width)
+            image = np.tile(gradient, (self.image_height, 1))
+            return np.stack((image,) * 3, axis=-1).astype(np.uint8)
+        elif self.bg_source == BGSource.TOP:
+            gradient = np.linspace(max_pix, min_pix, self.image_height)[:, None]
+            image = np.tile(gradient, (1, self.image_width))
+            return np.stack((image,) * 3, axis=-1).astype(np.uint8)
+        elif self.bg_source == BGSource.BOTTOM:
+            gradient = np.linspace(min_pix, max_pix, self.image_height)[:, None]
+            image = np.tile(gradient, (1, self.image_width))
+            return np.stack((image,) * 3, axis=-1).astype(np.uint8)
+        else:
+            raise ValueError('Wrong initial latent!')
+    
+    @torch.no_grad()
+    def __call__(self, input_video, init_latent=None, input_strength=None):
+        input_latent = encode_video(self.vae, input_video)* self.vae.config.scaling_factor
+        
+        if input_strength:
+            light_strength = input_strength
+        else:
+            light_strength = self.lowres_denoise
+
+        if not init_latent:
+            init_latent = self.bg_latent
+
+        latents = self.pipeline(
+            image=init_latent,
+            strength=light_strength,
+            prompt_embeds=self.conds,
+            negative_prompt_embeds=self.unconds,
+            width=self.image_width,
+            height=self.image_height,
+            num_inference_steps=int(round(self.steps / self.lowres_denoise)),
+            num_images_per_prompt=self.num_samples,
+            generator=self.generator,
+            output_type='latent',
+            guidance_scale=self.cfg,
+            cross_attention_kwargs={'concat_conds': input_latent},
+        ).images.to(self.pipeline.vae.dtype)
+
+        relight_video = decode_latents(self.vae, latents)
+        return relight_video

src/tools.py

@@ -0,0 +1,152 @@
+from PIL import Image,ImageSequence
+import numpy as np
+import torch
+try:
+    from moviepy import VideoFileClip
+except:
+    from moviepy.editor import VideoFileClip
+
+import os
+import imageio
+import random
+from diffusers.utils import  export_to_video
+
+def resize_and_center_crop(image, target_width, target_height):
+    pil_image = Image.fromarray(image)
+    original_width, original_height = pil_image.size
+    scale_factor = max(target_width / original_width, target_height / original_height)
+    resized_width = int(round(original_width * scale_factor))
+    resized_height = int(round(original_height * scale_factor))
+    resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
+
+    left = (resized_width - target_width) / 2
+    top = (resized_height - target_height) / 2
+    right = (resized_width + target_width) / 2
+    bottom = (resized_height + target_height) / 2
+    cropped_image = resized_image.crop((left, top, right, bottom))
+    return np.array(cropped_image)
+
+def numpy2pytorch(imgs, device, dtype):
+    h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.0 - 1.0
+    h = h.movedim(-1, 1)
+    return h.to(device=device, dtype=dtype)
+
+def get_fg_video(video_list, mask_list, device, dtype):
+    video_np = np.stack(video_list, axis=0)
+    mask_np = np.stack(mask_list, axis=0)
+    mask_bool = mask_np == 255
+    video_fg = np.where(mask_bool, video_np, 127)
+
+    h = torch.from_numpy(video_fg).float() / 127.0 - 1.0
+    h = h.movedim(-1, 1)
+    return h.to(device=device, dtype=dtype)
+
+
+def pad(x, p, i):
+    return x[:i] if len(x) >= i else x + [p] * (i - len(x))
+
+def gif_to_mp4(gif_path, mp4_path):
+    clip = VideoFileClip(gif_path)
+    clip.write_videofile(mp4_path)
+
+def generate_light_sequence(light_tensor, num_frames=16, direction="r"):
+
+    if direction in "l":
+        target_tensor = torch.rot90(light_tensor, k=1, dims=(2, 3))
+    elif direction in "r":
+        target_tensor = torch.rot90(light_tensor, k=-1, dims=(2, 3))
+    else:
+        raise ValueError("direction must be either 'r' for right or 'l' for left")
+        
+    # Generate the sequence
+    out_list = []
+    for frame_idx in range(num_frames):
+        t = frame_idx / (num_frames - 1) 
+        interpolated_matrix = (1 - t) * light_tensor + t * target_tensor
+        out_list.append(interpolated_matrix)
+    
+    out_tensor = torch.stack(out_list, dim=0).squeeze(1)
+
+    return out_tensor
+
+def tensor2vid(video: torch.Tensor, processor, output_type="np"):
+
+    batch_size, channels, num_frames, height, width = video.shape ## [1, 4, 16, 512, 512]
+    outputs = []
+    for batch_idx in range(batch_size):
+        batch_vid = video[batch_idx].permute(1, 0, 2, 3)
+        batch_output = processor.postprocess(batch_vid, output_type)
+
+        outputs.append(batch_output)
+
+    return outputs
+
+def read_video(video_path:str, image_width, image_height):
+    extension = video_path.split('.')[-1].lower()
+    video_name = os.path.basename(video_path)
+    video_list = []
+
+    if extension in "gif":
+        ## input from gif
+        video = Image.open(video_path)
+        for i, frame in enumerate(ImageSequence.Iterator(video)):
+            frame = np.array(frame.convert("RGB"))
+            frame = resize_and_center_crop(frame, image_width, image_height)
+            video_list.append(frame)
+    elif extension in "mp4":
+        ## input from mp4
+        reader = imageio.get_reader(video_path)
+        for frame in reader:
+            frame = resize_and_center_crop(frame, image_width, image_height)
+            video_list.append(frame)
+    else:
+        raise ValueError('Wrong input type')
+    
+    video_list = [Image.fromarray(frame) for frame in video_list]
+
+    return video_list, video_name
+
+def read_mask(mask_folder:str):
+    mask_files = os.listdir(mask_folder)
+    mask_files = sorted(mask_files)
+    mask_list = []
+    for mask_file in mask_files:
+        mask_path = os.path.join(mask_folder, mask_file)
+        mask = Image.open(mask_path).convert('RGB')
+        mask_list.append(mask)
+    
+    return mask_list
+
+def decode_latents(vae, latents, decode_chunk_size: int = 16):
+    
+    latents = 1 / vae.config.scaling_factor * latents
+    video = []
+    for i in range(0, latents.shape[0], decode_chunk_size):
+        batch_latents = latents[i : i + decode_chunk_size]
+        batch_latents = vae.decode(batch_latents).sample
+        video.append(batch_latents)
+
+    video = torch.cat(video)
+
+    return video
+
+def encode_video(vae, video, decode_chunk_size: int = 16) -> torch.Tensor:
+    latents = []
+    for i in range(0, len(video), decode_chunk_size):
+        batch_video = video[i : i + decode_chunk_size]
+        batch_video = vae.encode(batch_video).latent_dist.mode()
+        latents.append(batch_video)
+    return torch.cat(latents)
+
+def vis_video(input_video, video_processor, save_path):
+    ## shape: 1, c, f, h, w
+    relight_video = video_processor.postprocess_video(video=input_video, output_type="pil")
+    export_to_video(relight_video[0], save_path)
+    
+def set_all_seed(seed):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True