这是一个基于stable diffusion的扩展绘画工具（outpainting）
from diffusers import UNet2DConditionModel
from diffusers import AutoencoderKL
from diffusers import DDIMScheduler
from diffusers.image_processor import VaeImageProcessor
from transformers import CLIPTokenizer, CLIPTextModel
import torch
from tqdm import tqdm
import cv2 as cv
import numpy as np
import time
from safetensors.torch import load_file
from typing import List



class ImageExtension():
    DEFAULT_MODEL = "stabilityai/stable-diffusion-2-inpainting"
    cache_dir = "F:/huggingface_model/"    # diffusers的本地缓存路径
    def __init__(self, 
                 sampleTimeStep:int = 30, 
                 only_local_files:bool = False):
        self.base_model_name = self.DEFAULT_MODEL
        self.only_local_files = only_local_files
        self.sampleTimeStep = sampleTimeStep
        self.load_model()
        self.getLatent_model()
        self.load_text_encoder()
        self.load_scheduler()
        
        self.allTimestep = self.ddim.timesteps

        self.image_processor = VaeImageProcessor()


    def addNoise(self, latent:torch.Tensor, noise: torch.Tensor, timestep:torch.Tensor):
        latent = self.ddim.add_noise(latent, noise, timestep)
        latent = latent * self.ddim.init_noise_sigma
    
        return latent
    
    def set_step(self, sampleTimeStep:int):
        self.sampleTimeStep = sampleTimeStep
        self.ddim.set_timesteps(self.sampleTimeStep, device="cuda:0")
        self.allTimestep = self.ddim.timesteps
    
        
    def sample_step(self, latent: torch.Tensor, niose: torch.Tensor, timestep: torch.Tensor):
        return self.ddim.step(niose, timestep, latent)['prev_sample']


    def sample_block(self, latent:torch.Tensor, masked_latent:torch.Tensor, mask: torch.Tensor, prompt_embeds:torch.Tensor, timestep: torch.Tensor, guidance_scale:int=7):
        latent_model_input = torch.cat([latent] * 2)
        mask_input = torch.cat([mask]*2)
        masked_latent_input = torch.cat([masked_latent]*2)

        latent_model_input = self.ddim.scale_model_input(latent_model_input, timestep)
        latent_model_input = torch.cat([latent_model_input, mask_input, masked_latent_input], dim=1)     # inpaint模型拥有额外的输入信息，通道数为9
        
        noise_pred = self.unet(latent_model_input, timestep, encoder_hidden_states = prompt_embeds).sample
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

        latent = self.sample_step(latent, noise_pred, timestep)

        return latent
    
    def sample(self, latent:torch.Tensor, masked_latent:torch.Tensor, mask: torch.Tensor, chunk_list: List[int], prompt_embeds:torch.Tensor, guidance_scale:int=7):
        # print(prompt_embeds.shape)
        # print(latent.shape)
        count = torch.zeros_like(latent)
        full_latent = torch.zeros_like(latent)
        for Tin in tqdm(range(0, len(self.allTimestep))):
            Ti = self.allTimestep[Tin]
            count.zero_()
            full_latent.zero_()
            for chunk_block in chunk_list:

                sample_latent = latent[:, :, chunk_block[0]:chunk_block[1], chunk_block[2]:chunk_block[3]]
                sample_mask = mask[:, :, chunk_block[0]:chunk_block[1], chunk_block[2]:chunk_block[3]]
                sample_masked_latent = masked_latent[:, :, chunk_block[0]:chunk_block[1], chunk_block[2]:chunk_block[3]]

                pred_noise = self.sample_block(sample_latent, sample_masked_latent, sample_mask, prompt_embeds, Ti, guidance_scale)   # 每一个时间步的采样过程
                
                full_latent[:, :, chunk_block[0]:chunk_block[1], chunk_block[2]:chunk_block[3]] += pred_noise
                count[:, :, chunk_block[0]:chunk_block[1], chunk_block[2]:chunk_block[3]] += 1


            latent = torch.where(count > 0, full_latent / count, full_latent)

        return latent
    
    def load_scheduler(self):
        MN = self.base_model_name
        self.ddim = DDIMScheduler.from_pretrained(MN, 
                                                  subfolder="scheduler", 
                                                  local_files_only=self.only_local_files, 
                                                #   torch_dtype=torch.float16, 
                                                  use_safetensors=True, 
                                                  cache_dir = self.cache_dir)
        

        self.ddim.set_timesteps(self.sampleTimeStep, device="cuda:0")

    def load_model(self):
        self.unet = UNet2DConditionModel.from_pretrained(self.base_model_name, 
                                                            local_files_only = self.only_local_files, 
                                                            torch_dtype=torch.float16, 
                                                            # use_safetensors=True, 
                                                            subfolder = "unet",
                                                            cache_dir = self.cache_dir).cuda()
        
     
        self.unet.enable_xformers_memory_efficient_attention()
        
    def getLatent_model(self):
        MN = self.base_model_name
        self.vae = AutoencoderKL.from_pretrained(MN, 
                                                 local_files_only = self.only_local_files,
                                                 torch_dtype=torch.float16,
                                                #  use_safetensors=True,
                                                 subfolder = "vae",
                                                 cache_dir = self.cache_dir).cuda()
        

    def load_text_encoder(self):
        MN = self.base_model_name
        self.text_encoder = CLIPTextModel.from_pretrained(MN, 
                                                          local_files_only = self.only_local_files,
                                                          torch_dtype=torch.float16,
                                                        #   use_safetensors=True,
                                                          subfolder = "text_encoder",
                                                          cache_dir = self.cache_dir).cuda()
    
        self.tokenizer = CLIPTokenizer.from_pretrained(MN,
                                                         local_files_only = self.only_local_files,
                                                         subfolder = "tokenizer",
                                                         cache_dir = self.cache_dir)
        
        
    @staticmethod
    def tokenize_prompt(tokenizer, prompt):
        text_inputs = tokenizer(
            prompt,
            padding="max_length",
            max_length=tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids
        return text_input_ids 

    
    def encode_prompt(self, prompt:str, neg_prompt:str = None):
        text_input_ids = self.tokenize_prompt(self.tokenizer, prompt)

        prompt_embeds = self.text_encoder(
            text_input_ids.to(self.text_encoder.device),
            output_hidden_states=True,
        )
        encoder_hidden_states = prompt_embeds.hidden_states[-2]
        prompt_embeds = self.text_encoder.text_model.final_layer_norm(encoder_hidden_states)
        # prompt_embeds = prompt_embeds[0]

        if neg_prompt is None:
            neg_prompt = ""
        negative_text_input_ids = self.tokenize_prompt(self.tokenizer, neg_prompt)
        negative_prompt_embeds = self.text_encoder(
            negative_text_input_ids.to(self.text_encoder.device),
            output_hidden_states=True,
        )
        negative_prompt_embeds = negative_prompt_embeds[0]
    
        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])

        return prompt_embeds

    def get_text_embedding(self, prompt:str, neg_prompt:str = None):
        return self.encode_prompt(prompt, neg_prompt)

        
    def getImgLatent(self, img:torch.Tensor):
        # img = self.image_processor.preprocess(img)
        return self.vae.encode(img).latent_dist.sample() * self.vae.config.scaling_factor
    
    def getImg(self, latent:torch.Tensor):
        image = self.vae.decode(latent / self.vae.config.scaling_factor)[0]
        image = image.detach()
        image = self.image_processor.postprocess(image, output_type="np", do_denormalize=[True])
        return image

    
def main():
    pass

if __name__ == "__main__":
    main()


imageExtension.py
WatchTower-Liu/diffusion-outpainting
import gradio as gr
import numpy as np
import cv2 as cv
import torch

from imageExtension import ImageExtension

class Event():
    def __init__(self) -> None:
        self.col_move = 0
        self.row_move = 0
        self.grid_size = 128

        self.img_ext = ImageExtension(only_local_files=False)

    def set_grid_size(self, grid_size:int):
        self.grid_size = int(grid_size)

    def image_change(self, img:np.ndarray):
        selection = [([0, 0, img.shape[1], img.shape[0]], "image")]
        return img, selection
    
    def render_box(self, img:np.ndarray):
        imgShape = img.shape
        selection = [([max(-self.col_move, 0) * self.grid_size, 
                       max(-self.row_move, 0) * self.grid_size, 
                       max(-self.col_move, 0) * self.grid_size + imgShape[1], 
                       max(-self.row_move, 0) * self.grid_size + imgShape[0]], "image")]
        if self.col_move != 0 or self. row_move != 0:
            selection.append(
                     ([max(self.col_move, 0) * self.grid_size, 
                       max(self.row_move, 0) * self.grid_size, 
                       max(self.col_move, 0) * self.grid_size + imgShape[1], 
                       max(self.row_move, 0) * self.grid_size + imgShape[0]], "extension")
            )

        newImgShape = [
            abs(self.row_move) * self.grid_size + imgShape[0],
            abs(self.col_move) * self.grid_size + imgShape[1]
        ]
        newImg = np.zeros([*newImgShape, 3], dtype=np.uint8)
        newImg[selection[0][0][1] : selection[0][0][3], selection[0][0][0] : selection[0][0][2]] = img

        return (newImg, selection)
    
    def left_ext(self, img:np.ndarray):
        imgShape = img.shape
        
        self.col_move -= 1
        self.col_move = self.col_move if abs(self.col_move * self.grid_size) < imgShape[1]//2 else self.col_move + 1

        return self.render_box(img)
    
    def right_ext(self, img:np.ndarray):
        imgShape = img.shape
        
        self.col_move += 1
        self.col_move = self.col_move if self.col_move * self.grid_size < imgShape[1]//2 else self.col_move - 1

        return self.render_box(img)
    
    def up_ext(self, img:np.ndarray):
        imgShape = img.shape
        
        self.row_move -= 1
        self.row_move = self.row_move if abs(self.row_move * self.grid_size) < imgShape[0]//2 else self.row_move + 1

        return self.render_box(img)
    
    def down_ext(self, img:np.ndarray):
        imgShape = img.shape
        
        self.row_move += 1
        self.row_move = self.row_move if self.row_move * self.grid_size < imgShape[0]//2 else self.row_move - 1

        return self.render_box(img)
    
    def process(self, img:np.ndarray, prompt:str, neg_prompt:str, sampleStep:int, guidance_scale:int):
        newimg, bbox = self.render_box(img)
        torch.cuda.empty_cache()

        self.img_ext.set_step(sampleStep)
        with torch.no_grad():

            torchImg = torch.from_numpy(newimg).permute(2, 0, 1).unsqueeze(0).half()
            torchUseImg = torchImg[:, :, bbox[0][0][1] : bbox[0][0][3], bbox[0][0][0] : bbox[0][0][2]] / 127.5 - 1.0

            torchImg[:, :, bbox[0][0][1] : bbox[0][0][3], bbox[0][0][0] : bbox[0][0][2]] = torchUseImg

            torchImg = torchImg.cuda().half()
            mask = np.zeros([torchImg.shape[2], torchImg.shape[3]], dtype=np.uint8)
            mask[bbox[1][0][1] : bbox[1][0][3], bbox[1][0][0] : bbox[1][0][2]] = 1      # 需要扩展的地方设置为mask 1
            mask[bbox[0][0][1] : bbox[0][0][3], bbox[0][0][0] : bbox[0][0][2]] = 0       # 重叠区域再次置零

            mask = cv.medianBlur(mask, 3)
            mask = mask[None, None, :, :]
            mask = torch.from_numpy(mask).cuda().half()

            masked_img = torchImg * (mask < 0.5)  # 反向掩码
            latent = self.img_ext.getImgLatent(torchImg)
            masked_latent = self.img_ext.getImgLatent(masked_img)

            mask = torch.nn.functional.interpolate(
                mask, size=(mask.shape[2] // 8, mask.shape[3] // 8)                            # mask 转换为latent尺寸
            )

            noise = torch.randn_like(latent)
            extend_latent = self.img_ext.addNoise(latent, noise, self.img_ext.allTimestep[0])  # 加噪

            chunkList = [
                [bbox[0][0][1]//8, bbox[0][0][3]//8, bbox[0][0][0]//8, bbox[0][0][2]//8], 
                [bbox[1][0][1]//8, bbox[1][0][3]//8, bbox[1][0][0]//8, bbox[1][0][2]//8]
                ]


            prompt_embeds = self.img_ext.get_text_embedding(prompt, neg_prompt)
        
            result = self.img_ext.sample(extend_latent, masked_latent, mask, chunkList, prompt_embeds, guidance_scale=guidance_scale)    
            out_img = self.img_ext.getImg(result)[0]

        out_img = cv.resize(out_img, (newimg.shape[1], newimg.shape[0]))

        outPaintMask = np.zeros_like(out_img, dtype=np.uint8)
        outPaintMask[bbox[1][0][1] : bbox[1][0][3], bbox[1][0][0] : bbox[1][0][2]] = 1   
        outPaintMask[bbox[0][0][1] : bbox[0][0][3], bbox[0][0][0] : bbox[0][0][2]] = 1

        out_img = np.where(outPaintMask < 0.5, newimg, out_img)   

        return out_img
    
    def start_process(self, img:np.ndarray, prompt:str, neg_prompt:str, sampleStep:int, guidance_scale:int):
        return self.process(img, prompt, neg_prompt, sampleStep, guidance_scale)
    
    def apply_to(self, img):
        self.col_move = 0
        self.row_move = 0

        return img


def setupUI(event_process:Event):
    with gr.Blocks() as UI:
        with gr.Row():
            with gr.Column():
                image = gr.Image(height=300)
                img_window = gr.AnnotatedImage(
                    color_map={"image": "#a89a00", "extension": "#00aeff"},
                    height = 300
                )
                image.upload(event_process.image_change, inputs=[image], outputs=[img_window])
                image.clear(event_process.apply_to, inputs=[image])
            pre_image = gr.Image(height=600, interactive=False)

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(interactive=True, label="prompt")
                neg_prompt = gr.Textbox(interactive=True, value="Blurred, incomplete, distorted, misplaced", label="negative prompt")
                guidance_scale = gr.Slider(0, 20, value=7, label="guidance scale", step=0.5)
            with gr.Column():
                with gr.Row():
                    left = gr.Button("LEFT")
                    with gr.Column():
                        up = gr.Button("UP")
                        down = gr.Button("DOWN")
                    right = gr.Button("right")

                    left.click(event_process.left_ext, inputs=[image], outputs=[img_window])
                    right.click(event_process.right_ext, inputs=[image], outputs=[img_window])
                    up.click(event_process.up_ext, inputs=[image], outputs=[img_window])
                    down.click(event_process.down_ext, inputs=[image], outputs=[img_window])
                grid_size = gr.Number(128, label="grid size")
                grid_size.change(event_process.set_grid_size, inputs=[grid_size])
        with gr.Row():
            with gr.Row():
                process = gr.Button("start")
                apply = gr.Button("apply")
            sample_step = gr.Slider(0, 100, value=30, label="sample step", step=1)
            process.click(event_process.start_process, inputs=[image, prompt, neg_prompt, sample_step, guidance_scale], outputs=[pre_image])
            apply.click(event_process.apply_to, inputs=[pre_image], outputs=[image])


    return UI


def main():
    event_process = Event()
    UI = setupUI(event_process)
    UI.launch()

if __name__ == "__main__":
    main() 


webUI.py
WatchTower-Liu/diffusion-outpainting