modelscope/DiffSynth-Painter

1from ..models import SDTextEncoder, SDUNet, SDVAEDecoder, SDVAEEncoder, SDIpAdapter, IpAdapterCLIPImageEmbedder, SDMotionModel2from ..models.model_manager import ModelManager3from ..controlnets import MultiControlNetManager, ControlNetUnit, ControlNetConfigUnit, Annotator4from ..prompters import SDPrompter5from ..schedulers import EnhancedDDIMScheduler6from .sd_image import SDImagePipeline7from .dancer import lets_dance8from typing import List9import torch10from tqdm import tqdm11 12 13 14def lets_dance_with_long_video(15    unet: SDUNet,16    motion_modules: SDMotionModel = None,17    controlnet: MultiControlNetManager = None,18    sample = None,19    timestep = None,20    encoder_hidden_states = None,21    ipadapter_kwargs_list = {},22    controlnet_frames = None,23    unet_batch_size = 1,24    controlnet_batch_size = 1,25    cross_frame_attention = False,26    tiled=False,27    tile_size=64,28    tile_stride=32,29    device="cuda",30    animatediff_batch_size=16,31    animatediff_stride=8,32):33    num_frames = sample.shape[0]34    hidden_states_output = [(torch.zeros(sample[0].shape, dtype=sample[0].dtype), 0) for i in range(num_frames)]35 36    for batch_id in range(0, num_frames, animatediff_stride):37        batch_id_ = min(batch_id + animatediff_batch_size, num_frames)38 39        # process this batch40        hidden_states_batch = lets_dance(41            unet, motion_modules, controlnet,42            sample[batch_id: batch_id_].to(device),43            timestep,44            encoder_hidden_states,45            ipadapter_kwargs_list=ipadapter_kwargs_list,46            controlnet_frames=controlnet_frames[:, batch_id: batch_id_].to(device) if controlnet_frames is not None else None,47            unet_batch_size=unet_batch_size, controlnet_batch_size=controlnet_batch_size,48            cross_frame_attention=cross_frame_attention,49            tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, device=device50        ).cpu()51 52        # update hidden_states53        for i, hidden_states_updated in zip(range(batch_id, batch_id_), hidden_states_batch):54            bias = max(1 - abs(i - (batch_id + batch_id_ - 1) / 2) / ((batch_id_ - batch_id - 1 + 1e-2) / 2), 1e-2)55            hidden_states, num = hidden_states_output[i]56            hidden_states = hidden_states * (num / (num + bias)) + hidden_states_updated * (bias / (num + bias))57            hidden_states_output[i] = (hidden_states, num + bias)58 59        if batch_id_ == num_frames:60            break61 62    # output63    hidden_states = torch.stack([h for h, _ in hidden_states_output])64    return hidden_states65 66 67 68class SDVideoPipeline(SDImagePipeline):69 70    def __init__(self, device="cuda", torch_dtype=torch.float16, use_original_animatediff=True):71        super().__init__(device=device, torch_dtype=torch_dtype)72        self.scheduler = EnhancedDDIMScheduler(beta_schedule="linear" if use_original_animatediff else "scaled_linear")73        self.prompter = SDPrompter()74        # models75        self.text_encoder: SDTextEncoder = None76        self.unet: SDUNet = None77        self.vae_decoder: SDVAEDecoder = None78        self.vae_encoder: SDVAEEncoder = None79        self.controlnet: MultiControlNetManager = None80        self.ipadapter_image_encoder: IpAdapterCLIPImageEmbedder = None81        self.ipadapter: SDIpAdapter = None82        self.motion_modules: SDMotionModel = None83 84 85    def fetch_models(self, model_manager: ModelManager, controlnet_config_units: List[ControlNetConfigUnit]=[], prompt_refiner_classes=[]):86        # Main models87        self.text_encoder = model_manager.fetch_model("sd_text_encoder")88        self.unet = model_manager.fetch_model("sd_unet")89        self.vae_decoder = model_manager.fetch_model("sd_vae_decoder")90        self.vae_encoder = model_manager.fetch_model("sd_vae_encoder")91        self.prompter.fetch_models(self.text_encoder)92        self.prompter.load_prompt_refiners(model_manager, prompt_refiner_classes)93 94        # ControlNets95        controlnet_units = []96        for config in controlnet_config_units:97            controlnet_unit = ControlNetUnit(98                Annotator(config.processor_id, device=self.device),99                model_manager.fetch_model("sd_controlnet", config.model_path),100                config.scale101            )102            controlnet_units.append(controlnet_unit)103        self.controlnet = MultiControlNetManager(controlnet_units)104 105        # IP-Adapters106        self.ipadapter = model_manager.fetch_model("sd_ipadapter")107        self.ipadapter_image_encoder = model_manager.fetch_model("sd_ipadapter_clip_image_encoder")108 109        # Motion Modules110        self.motion_modules = model_manager.fetch_model("sd_motion_modules")111        if self.motion_modules is None:112            self.scheduler = EnhancedDDIMScheduler(beta_schedule="scaled_linear")113 114 115    @staticmethod116    def from_model_manager(model_manager: ModelManager, controlnet_config_units: List[ControlNetConfigUnit]=[], prompt_refiner_classes=[]):117        pipe = SDVideoPipeline(118            device=model_manager.device,119            torch_dtype=model_manager.torch_dtype,120        )121        pipe.fetch_models(model_manager, controlnet_config_units, prompt_refiner_classes)122        return pipe123    124 125    def decode_video(self, latents, tiled=False, tile_size=64, tile_stride=32):126        images = [127            self.decode_image(latents[frame_id: frame_id+1], tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)128            for frame_id in range(latents.shape[0])129        ]130        return images131    132 133    def encode_video(self, processed_images, tiled=False, tile_size=64, tile_stride=32):134        latents = []135        for image in processed_images:136            image = self.preprocess_image(image).to(device=self.device, dtype=self.torch_dtype)137            latent = self.encode_image(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)138            latents.append(latent.cpu())139        latents = torch.concat(latents, dim=0)140        return latents141    142 143    @torch.no_grad()144    def __call__(145        self,146        prompt,147        negative_prompt="",148        cfg_scale=7.5,149        clip_skip=1,150        num_frames=None,151        input_frames=None,152        ipadapter_images=None,153        ipadapter_scale=1.0,154        controlnet_frames=None,155        denoising_strength=1.0,156        height=512,157        width=512,158        num_inference_steps=20,159        animatediff_batch_size = 16,160        animatediff_stride = 8,161        unet_batch_size = 1,162        controlnet_batch_size = 1,163        cross_frame_attention = False,164        smoother=None,165        smoother_progress_ids=[],166        tiled=False,167        tile_size=64,168        tile_stride=32,169        progress_bar_cmd=tqdm,170        progress_bar_st=None,171    ):172        # Tiler parameters, batch size ...173        tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}174        other_kwargs = {175            "animatediff_batch_size": animatediff_batch_size, "animatediff_stride": animatediff_stride,176            "unet_batch_size": unet_batch_size, "controlnet_batch_size": controlnet_batch_size,177            "cross_frame_attention": cross_frame_attention,178        }179 180        # Prepare scheduler181        self.scheduler.set_timesteps(num_inference_steps, denoising_strength)182 183        # Prepare latent tensors184        if self.motion_modules is None:185            noise = torch.randn((1, 4, height//8, width//8), device="cpu", dtype=self.torch_dtype).repeat(num_frames, 1, 1, 1)186        else:187            noise = torch.randn((num_frames, 4, height//8, width//8), device="cpu", dtype=self.torch_dtype)188        if input_frames is None or denoising_strength == 1.0:189            latents = noise190        else:191            latents = self.encode_video(input_frames, **tiler_kwargs)192            latents = self.scheduler.add_noise(latents, noise, timestep=self.scheduler.timesteps[0])193 194        # Encode prompts195        prompt_emb_posi = self.encode_prompt(prompt, clip_skip=clip_skip, positive=True)196        prompt_emb_nega = self.encode_prompt(negative_prompt, clip_skip=clip_skip, positive=False)197 198        # IP-Adapter199        if ipadapter_images is not None:200            ipadapter_image_encoding = self.ipadapter_image_encoder(ipadapter_images)201            ipadapter_kwargs_list_posi = {"ipadapter_kwargs_list": self.ipadapter(ipadapter_image_encoding, scale=ipadapter_scale)}202            ipadapter_kwargs_list_nega = {"ipadapter_kwargs_list": self.ipadapter(torch.zeros_like(ipadapter_image_encoding))}203        else:204            ipadapter_kwargs_list_posi, ipadapter_kwargs_list_nega = {"ipadapter_kwargs_list": {}}, {"ipadapter_kwargs_list": {}}205 206        # Prepare ControlNets207        if controlnet_frames is not None:208            if isinstance(controlnet_frames[0], list):209                controlnet_frames_ = []210                for processor_id in range(len(controlnet_frames)):211                    controlnet_frames_.append(212                        torch.stack([213                            self.controlnet.process_image(controlnet_frame, processor_id=processor_id).to(self.torch_dtype)214                            for controlnet_frame in progress_bar_cmd(controlnet_frames[processor_id])215                        ], dim=1)216                    )217                controlnet_frames = torch.concat(controlnet_frames_, dim=0)218            else:219                controlnet_frames = torch.stack([220                    self.controlnet.process_image(controlnet_frame).to(self.torch_dtype)221                    for controlnet_frame in progress_bar_cmd(controlnet_frames)222                ], dim=1)223            controlnet_kwargs = {"controlnet_frames": controlnet_frames}224        else:225            controlnet_kwargs = {"controlnet_frames": None}226        227        # Denoise228        for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):229            timestep = timestep.unsqueeze(0).to(self.device)230 231            # Classifier-free guidance232            noise_pred_posi = lets_dance_with_long_video(233                self.unet, motion_modules=self.motion_modules, controlnet=self.controlnet,234                sample=latents, timestep=timestep,235                **prompt_emb_posi, **controlnet_kwargs, **ipadapter_kwargs_list_posi, **other_kwargs, **tiler_kwargs,236                device=self.device,237            )238            noise_pred_nega = lets_dance_with_long_video(239                self.unet, motion_modules=self.motion_modules, controlnet=self.controlnet,240                sample=latents, timestep=timestep,241                **prompt_emb_nega, **controlnet_kwargs, **ipadapter_kwargs_list_nega, **other_kwargs, **tiler_kwargs,242                device=self.device,243            )244            noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)245 246            # DDIM and smoother247            if smoother is not None and progress_id in smoother_progress_ids:248                rendered_frames = self.scheduler.step(noise_pred, timestep, latents, to_final=True)249                rendered_frames = self.decode_video(rendered_frames)250                rendered_frames = smoother(rendered_frames, original_frames=input_frames)251                target_latents = self.encode_video(rendered_frames)252                noise_pred = self.scheduler.return_to_timestep(timestep, latents, target_latents)253            latents = self.scheduler.step(noise_pred, timestep, latents)254 255            # UI256            if progress_bar_st is not None:257                progress_bar_st.progress(progress_id / len(self.scheduler.timesteps))258        259        # Decode image260        output_frames = self.decode_video(latents, **tiler_kwargs)261 262        # Post-process263        if smoother is not None and (num_inference_steps in smoother_progress_ids or -1 in smoother_progress_ids):264            output_frames = smoother(output_frames, original_frames=input_frames)265 266        return output_frames267