modelscope/DiffSynth-Painter

1import torch2import numpy as np3from PIL import Image4 5 6 7class BasePipeline(torch.nn.Module):8 9    def __init__(self, device="cuda", torch_dtype=torch.float16):10        super().__init__()11        self.device = device12        self.torch_dtype = torch_dtype13 14 15    def preprocess_image(self, image):16        image = torch.Tensor(np.array(image, dtype=np.float32) * (2 / 255) - 1).permute(2, 0, 1).unsqueeze(0)17        return image18    19 20    def preprocess_images(self, images):21        return [self.preprocess_image(image) for image in images]22    23 24    def vae_output_to_image(self, vae_output):25        image = vae_output[0].cpu().float().permute(1, 2, 0).numpy()26        image = Image.fromarray(((image / 2 + 0.5).clip(0, 1) * 255).astype("uint8"))27        return image28    29 30    def vae_output_to_video(self, vae_output):31        video = vae_output.cpu().permute(1, 2, 0).numpy()32        video = [Image.fromarray(((image / 2 + 0.5).clip(0, 1) * 255).astype("uint8")) for image in video]33        return video34 35    36    def merge_latents(self, value, latents, masks, scales):37        height, width = value.shape[-2:]38        weight = torch.ones_like(value)39        for latent, mask, scale in zip(latents, masks, scales):40            mask = self.preprocess_image(mask.resize((width, height))).mean(dim=1, keepdim=True) > 041            mask = mask.repeat(1, latent.shape[1], 1, 1)42            value[mask] += latent[mask] * scale43            weight[mask] += scale44        value /= weight45        return value46 47 48    def control_noise_via_local_prompts(self, prompt_emb_global, prompt_emb_locals, masks, mask_scales, inference_callback):49        noise_pred_global = inference_callback(prompt_emb_global)50        noise_pred_locals = [inference_callback(prompt_emb_local) for prompt_emb_local in prompt_emb_locals]51        noise_pred = self.merge_latents(noise_pred_global, noise_pred_locals, masks, mask_scales)52        return noise_pred53