modelscope/DiffSynth-Painter

1from ..models import ModelManager, SD3TextEncoder1, SD3TextEncoder2, SD3TextEncoder3, SD3DiT, SD3VAEDecoder, SD3VAEEncoder2from ..prompters import SD3Prompter3from ..schedulers import FlowMatchScheduler4from .base import BasePipeline5import torch6from tqdm import tqdm7 8 9 10class SD3ImagePipeline(BasePipeline):11 12    def __init__(self, device="cuda", torch_dtype=torch.float16):13        super().__init__(device=device, torch_dtype=torch_dtype)14        self.scheduler = FlowMatchScheduler()15        self.prompter = SD3Prompter()16        # models17        self.text_encoder_1: SD3TextEncoder1 = None18        self.text_encoder_2: SD3TextEncoder2 = None19        self.text_encoder_3: SD3TextEncoder3 = None20        self.dit: SD3DiT = None21        self.vae_decoder: SD3VAEDecoder = None22        self.vae_encoder: SD3VAEEncoder = None23 24 25    def denoising_model(self):26        return self.dit27 28 29    def fetch_models(self, model_manager: ModelManager, prompt_refiner_classes=[]):30        self.text_encoder_1 = model_manager.fetch_model("sd3_text_encoder_1")31        self.text_encoder_2 = model_manager.fetch_model("sd3_text_encoder_2")32        self.text_encoder_3 = model_manager.fetch_model("sd3_text_encoder_3")33        self.dit = model_manager.fetch_model("sd3_dit")34        self.vae_decoder = model_manager.fetch_model("sd3_vae_decoder")35        self.vae_encoder = model_manager.fetch_model("sd3_vae_encoder")36        self.prompter.fetch_models(self.text_encoder_1, self.text_encoder_2, self.text_encoder_3)37        self.prompter.load_prompt_refiners(model_manager, prompt_refiner_classes)38 39 40    @staticmethod41    def from_model_manager(model_manager: ModelManager, prompt_refiner_classes=[]):42        pipe = SD3ImagePipeline(43            device=model_manager.device,44            torch_dtype=model_manager.torch_dtype,45        )46        pipe.fetch_models(model_manager, prompt_refiner_classes)47        return pipe48    49 50    def encode_image(self, image, tiled=False, tile_size=64, tile_stride=32):51        latents = self.vae_encoder(image, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)52        return latents53    54 55    def decode_image(self, latent, tiled=False, tile_size=64, tile_stride=32):56        image = self.vae_decoder(latent.to(self.device), tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)57        image = self.vae_output_to_image(image)58        return image59    60 61    def encode_prompt(self, prompt, positive=True):62        prompt_emb, pooled_prompt_emb = self.prompter.encode_prompt(63            prompt, device=self.device, positive=positive64        )65        return {"prompt_emb": prompt_emb, "pooled_prompt_emb": pooled_prompt_emb}66    67 68    def prepare_extra_input(self, latents=None):69        return {}70    71 72    @torch.no_grad()73    def __call__(74        self,75        prompt,76        local_prompts=[],77        masks=[],78        mask_scales=[],79        negative_prompt="",80        cfg_scale=7.5,81        input_image=None,82        denoising_strength=1.0,83        height=1024,84        width=1024,85        num_inference_steps=20,86        tiled=False,87        tile_size=128,88        tile_stride=64,89        progress_bar_cmd=tqdm,90        progress_bar_st=None,91    ):92        # Tiler parameters93        tiler_kwargs = {"tiled": tiled, "tile_size": tile_size, "tile_stride": tile_stride}94 95        # Prepare scheduler96        self.scheduler.set_timesteps(num_inference_steps, denoising_strength)97 98        # Prepare latent tensors99        if input_image is not None:100            image = self.preprocess_image(input_image).to(device=self.device, dtype=self.torch_dtype)101            latents = self.encode_image(image, **tiler_kwargs)102            noise = torch.randn((1, 16, height//8, width//8), device=self.device, dtype=self.torch_dtype)103            latents = self.scheduler.add_noise(latents, noise, timestep=self.scheduler.timesteps[0])104        else:105            latents = torch.randn((1, 16, height//8, width//8), device=self.device, dtype=self.torch_dtype)106 107        # Encode prompts108        prompt_emb_posi = self.encode_prompt(prompt, positive=True)109        prompt_emb_nega = self.encode_prompt(negative_prompt, positive=False)110        prompt_emb_locals = [self.encode_prompt(prompt_local) for prompt_local in local_prompts]111 112        # Denoise113        for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):114            timestep = timestep.unsqueeze(0).to(self.device)115 116            # Classifier-free guidance117            inference_callback = lambda prompt_emb_posi: self.dit(118                latents, timestep=timestep, **prompt_emb_posi, **tiler_kwargs,119            )120            noise_pred_posi = self.control_noise_via_local_prompts(prompt_emb_posi, prompt_emb_locals, masks, mask_scales, inference_callback)121            noise_pred_nega = self.dit(122                latents, timestep=timestep, **prompt_emb_nega, **tiler_kwargs,123            )124            noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega)125 126            # DDIM127            latents = self.scheduler.step(noise_pred, self.scheduler.timesteps[progress_id], latents)128 129            # UI130            if progress_bar_st is not None:131                progress_bar_st.progress(progress_id / len(self.scheduler.timesteps))132        133        # Decode image134        image = self.decode_image(latents, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)135 136        return image137