使用HuggingFace实现 DiffEdit论文的掩码引导语义图像编辑

 import torch, logging
 
 ## disable warnings
 logging.disable(logging.WARNING)  
 
 ## Imaging  library
 from PIL import Image
 from torchvision import transforms as tfms
 
 
 ## Basic libraries
 from fastdownload import FastDownload
 import numpy as np
 from tqdm.auto import tqdm
 import matplotlib.pyplot as plt
 %matplotlib inline
 from IPython.display import display
 import shutil
 import os
 
 ## For video display
 from IPython.display import HTML
 from base64 import b64encode
 
 
 ## Import the CLIP artifacts 
 from transformers import CLIPTextModel, CLIPTokenizer
 from diffusers import AutoencoderKL, UNet2DConditionModel, DDIMScheduler
 
 ## Helper functions
 
 def load_artifacts():
     '''
     A function to load all diffusion artifacts
     '''
     vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16).to("cuda")
     unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16).to("cuda")
     tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16)
     text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16).to("cuda")
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)    
     return vae, unet, tokenizer, text_encoder, scheduler
 
 def load_image(p):
     '''
     Function to load images from a defined path
     '''
     return Image.open(p).convert('RGB').resize((512,512))
 
 def pil_to_latents(image):
     '''
     Function to convert image to latents
     '''
     init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
     init_image = init_image.to(device="cuda", dtype=torch.float16) 
     init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
     return init_latent_dist
 
 def latents_to_pil(latents):
     '''
     Function to convert latents to images
     '''
     latents = (1 / 0.18215) * latents
     with torch.no_grad():
         image = vae.decode(latents).sample
     image = (image / 2 + 0.5).clamp(0, 1)
     image = image.detach().cpu().permute(0, 2, 3, 1).numpy()
     images = (image * 255).round().astype("uint8")
     pil_images = [Image.fromarray(image) for image in images]
     return pil_images
 
 def text_enc(prompts, maxlen=None):
     '''
     A function to take a texual promt and convert it into embeddings
     '''
     if maxlen is None: maxlen = tokenizer.model_max_length
     inp = tokenizer(prompts, padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") 
     return text_encoder(inp.input_ids.to("cuda"))[0].half()
 
 vae, unet, tokenizer, text_encoder, scheduler = load_artifacts()

 p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
 init_img = load_image(p)
 init_img

 def prompt_2_img_i2i(prompts, init_img, neg_prompts=None, g=7.5, seed=100, strength =0.8, steps=50, dim=512):
     """
     Diffusion process to convert prompt to image
     """
     # Converting textual prompts to embedding
     text = text_enc(prompts) 
     
     # Adding an unconditional prompt , helps in the generation process
     if not neg_prompts: uncond =  text_enc([""], text.shape[1])
     else: uncond =  text_enc(neg_prompt, text.shape[1])
     emb = torch.cat([uncond, text])
     
     # Setting the seed
     if seed: torch.manual_seed(seed)
     
     # Setting number of steps in scheduler
     scheduler.set_timesteps(steps)
     
     # Convert the seed image to latent
     init_latents = pil_to_latents(init_img)
     
     # Figuring initial time step based on strength
     init_timestep = int(steps * strength) 
     timesteps = scheduler.timesteps[-init_timestep]
     timesteps = torch.tensor([timesteps], device="cuda")
     
     # Adding noise to the latents 
     noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
     init_latents = scheduler.add_noise(init_latents, noise, timesteps)
     latents = init_latents
     
     # Computing the timestep to start the diffusion loop
     t_start = max(steps - init_timestep, 0)
     timesteps = scheduler.timesteps[t_start:].to("cuda")
     
     # Iterating through defined steps
     for i,ts in enumerate(tqdm(timesteps)):
         # We need to scale the i/p latents to match the variance
         inp = scheduler.scale_model_input(torch.cat([latents] * 2), ts)
         
         # Predicting noise residual using U-Net
         with torch.no_grad(): u,t = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(2)
             
         # Performing Guidance
         pred = u + g*(t-u)
 
         # Conditioning  the latents
         #latents = scheduler.step(pred, ts, latents).pred_original_sample
         latents = scheduler.step(pred, ts, latents).prev_sample
     
     # Returning the latent representation to output an array of 4x64x64
     return latents.detach().cpu()

 def create_mask(init_img, rp, qp, n=10, s=0.5):
     ## Initialize a dictionary to save n iterations
     diff = {}
     
     ## Repeating the difference process n times
     for idx in range(n):
         ## Creating denoised sample using reference / original text
         orig_noise = prompt_2_img_i2i(prompts=rp, init_img=init_img, strength=s, seed = 100*idx)[0]
         ## Creating denoised sample using query / target text
         query_noise = prompt_2_img_i2i(prompts=qp, init_img=init_img, strength=s, seed = 100*idx)[0]
         ## Taking the difference 
         diff[idx] = (np.array(orig_noise)-np.array(query_noise))
     
     ## Creating a mask placeholder
     mask = np.zeros_like(diff[0])
     
     ## Taking an average of 10 iterations
     for idx in range(n):
         ## Note np.abs is a key step
         mask += np.abs(diff[idx])  
         
     ## Averaging multiple channels 
     mask = mask.mean(0)
     
     ## Normalizing 
     mask = (mask - mask.mean()) / np.std(mask)
     
     ## Binarizing and returning the mask object
     return (mask > 0).astype("uint8")
 
 mask = create_mask(init_img=init_img, rp=["a horse image"], qp=["a zebra image"], n=10)

 plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
 plt.imshow(
     Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
     cmap='cividis', 
     alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
 )

 def prompt_2_img_diffedit(rp, qp, init_img, mask, g=7.5, seed=100, strength =0.7, steps=70, dim=512):
     """
     Diffusion process to convert prompt to image
     """
     # Converting textual prompts to embedding
     rtext = text_enc(rp) 
     qtext = text_enc(qp)
     
     # Adding an unconditional prompt , helps in the generation process
     uncond =  text_enc([""], rtext.shape[1])
     emb = torch.cat([uncond, rtext, qtext])
     
     # Setting the seed
     if seed: torch.manual_seed(seed)
     
     # Setting number of steps in scheduler
     scheduler.set_timesteps(steps)
     
     # Convert the seed image to latent
     init_latents = pil_to_latents(init_img)
     
     # Figuring initial time step based on strength
     init_timestep = int(steps * strength) 
     timesteps = scheduler.timesteps[-init_timestep]
     timesteps = torch.tensor([timesteps], device="cuda")
     
     # Adding noise to the latents 
     noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
     init_latents = scheduler.add_noise(init_latents, noise, timesteps)
     latents = init_latents
     
     # Computing the timestep to start the diffusion loop
     t_start = max(steps - init_timestep, 0)
     timesteps = scheduler.timesteps[t_start:].to("cuda")
     
     # Converting mask to torch tensor
     mask = torch.tensor(mask, dtype=unet.dtype).unsqueeze(0).unsqueeze(0).to("cuda")
     
     # Iterating through defined steps
     for i,ts in enumerate(tqdm(timesteps)):
         # We need to scale the i/p latents to match the variance
         inp = scheduler.scale_model_input(torch.cat([latents] * 3), ts)
         
         # Predicting noise residual using U-Net
         with torch.no_grad(): u, rt, qt = unet(inp, ts, encoder_hidden_states=emb).sample.chunk(3)
             
         # Performing Guidance
         rpred = u + g*(rt-u)
         qpred = u + g*(qt-u)
 
         # Conditioning  the latents
         rlatents = scheduler.step(rpred, ts, latents).prev_sample
         qlatents = scheduler.step(qpred, ts, latents).prev_sample
         latents = mask*qlatents + (1-mask)*rlatents
     
     # Returning the latent representation to output an array of 4x64x64
     return latents_to_pil(latents)

 output = prompt_2_img_diffedit(
     rp = ["a horse image"], 
     qp=["a zebra image"],
     init_img=init_img, 
     mask = mask, 
     g=7.5, seed=100, strength =0.5, steps=70, dim=512)
 
 ## Plotting side by side
 fig, axs = plt.subplots(1, 2, figsize=(12, 6))
 for c, img in enumerate([init_img, output[0]]): 
     axs[c].imshow(img)
     if c == 0 : axs[c].set_title(f"Initial image ")
     else: axs[c].set_title(f"DiffEdit output")

 def diffEdit(init_img, rp , qp, g=7.5, seed=100, strength =0.7, steps=70, dim=512):
     
     ## Step 1: Create mask
     mask = create_mask(init_img=init_img, rp=rp, qp=qp)
     
     ## Step 2 and 3: Diffusion process using mask
     output = prompt_2_img_diffedit(
         rp = rp, 
         qp=qp, 
         init_img=init_img, 
         mask = mask, 
         g=g, 
         seed=seed,
         strength =strength, 
         steps=steps, 
         dim=dim)
     return mask , output

 def plot_diffEdit(init_img, output, mask):
     ## Plotting side by side
     fig, axs = plt.subplots(1, 3, figsize=(12, 6))
     
     ## Visualizing initial image
     axs[0].imshow(init_img)
     axs[0].set_title(f"Initial image")
     
     ## Visualizing initial image
     axs[2].imshow(output[0])
     axs[2].set_title(f"DiffEdit output")
     
     ## Visualizing the mask 
     axs[1].imshow(np.array(init_img), cmap='gray') 
     axs[1].imshow(
         Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
         cmap='cividis', 
         alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
     )
     axs[1].set_title(f"DiffEdit mask")

 p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
 init_img = load_image(p)
 mask, output = diffEdit(
   init_img, 
   rp = ["a horse image"], 
   qp=["a zebra image"]
 )
 plot_diffEdit(init_img, output, mask)

 p = FastDownload().download('https://raw.githubusercontent.com/johnrobinsn/diffusion_experiments/main/images/bowloberries_scaled.jpg')
 init_img = load_image(p)
 mask, output = diffEdit(
   init_img, 
   rp = ['Bowl of Strawberries'], 
   qp=['Bowl of Grapes']
 )
 plot_diffEdit(init_img, output, mask)

 def prompt_2_img_i2i_fast(prompts, init_img, g=7.5, seed=100, strength =0.5, steps=50, dim=512):
     """
     Diffusion process to convert prompt to image
     """
     # Converting textual prompts to embedding
     text = text_enc(prompts) 
     
     # Adding an unconditional prompt , helps in the generation process
     uncond =  text_enc([""], text.shape[1])
     emb = torch.cat([uncond, text])
     
     # Setting the seed
     if seed: torch.manual_seed(seed)
     
     # Setting number of steps in scheduler
     scheduler.set_timesteps(steps)
     
     # Convert the seed image to latent
     init_latents = pil_to_latents(init_img)
     
     # Figuring initial time step based on strength
     init_timestep = int(steps * strength) 
     timesteps = scheduler.timesteps[-init_timestep]
     timesteps = torch.tensor([timesteps], device="cuda")
     
     # Adding noise to the latents 
     noise = torch.randn(init_latents.shape, generator=None, device="cuda", dtype=init_latents.dtype)
     init_latents = scheduler.add_noise(init_latents, noise, timesteps)
     latents = init_latents
     
     # We need to scale the i/p latents to match the variance
     inp = scheduler.scale_model_input(torch.cat([latents] * 2), timesteps)
     # Predicting noise residual using U-Net
     with torch.no_grad(): u,t = unet(inp, timesteps, encoder_hidden_states=emb).sample.chunk(2)
          
     # Performing Guidance
     pred = u + g*(t-u)
 
     # Zero shot prediction
     latents = scheduler.step(pred, timesteps, latents).pred_original_sample
     
     # Returning the latent representation to output an array of 4x64x64
     return latents.detach().cpu()

 def create_mask_fast(init_img, rp, qp, n=20, s=0.5):
     ## Initialize a dictionary to save n iterations
     diff = {}
     
     ## Repeating the difference process n times
     for idx in range(n):
         ## Creating denoised sample using reference / original text
         orig_noise = prompt_2_img_i2i_fast(prompts=rp, init_img=init_img, strength=s, seed = 100*idx)[0]
         ## Creating denoised sample using query / target text
         query_noise = prompt_2_img_i2i_fast(prompts=qp, init_img=init_img, strength=s, seed = 100*idx)[0]
         ## Taking the difference 
         diff[idx] = (np.array(orig_noise)-np.array(query_noise))
     
     ## Creating a mask placeholder
     mask = np.zeros_like(diff[0])
     
     ## Taking an average of 10 iterations
     for idx in range(n):
         ## Note np.abs is a key step
         mask += np.abs(diff[idx])  
         
     ## Averaging multiple channels 
     mask = mask.mean(0)
     
     ## Normalizing 
     mask = (mask - mask.mean()) / np.std(mask)
     
     ## Binarizing and returning the mask object
     return (mask > 0).astype("uint8")

 p = FastDownload().download('https://images.pexels.com/photos/1996333/pexels-photo-1996333.jpeg?cs=srgb&dl=pexels-helena-lopes-1996333.jpg&fm=jpg&_gl=1*1pc0nw8*_ga*OTk4MTI0MzE4LjE2NjY1NDQwMjE.*_ga_8JE65Q40S6*MTY2Njc1MjIwMC4yLjEuMTY2Njc1MjIwMS4wLjAuMA..')
 init_img = load_image(p)
 mask = create_mask_fast(init_img=init_img, rp=["a horse image"], qp=["a zebra image"], n=20)
 plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
 plt.imshow(
     Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
     cmap='cividis', 
     alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
 )

 import cv2
 def improve_mask(mask):
     mask  = cv2.GaussianBlur(mask*255,(3,3),1) > 0
     return mask.astype('uint8')
 
 mask = improve_mask(mask)
 plt.imshow(np.array(init_img), cmap='gray') # I would add interpolation='none'
 plt.imshow(
     Image.fromarray(mask).resize((512,512)), ## Scaling the mask to original size
     cmap='cividis', 
     alpha=0.5*(np.array(Image.fromarray(mask*255).resize((512,512))) > 0)  
 )

 from diffusers import StableDiffusionInpaintPipeline
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
     revision="fp16",
     torch_dtype=torch.float16,
 ).to("cuda")

 pipe(
     prompt=["a zebra image"], 
     image=init_img, 
     mask_image=Image.fromarray(mask*255).resize((512,512)), 
     generator=torch.Generator("cuda").manual_seed(100),
     num_inference_steps = 20
 ).images[0]
 image

 def fastDiffEdit(init_img, rp , qp, g=7.5, seed=100, strength =0.7, steps=20, dim=512):
     
     ## Step 1: Create mask
     mask = create_mask_fast(init_img=init_img, rp=rp, qp=qp, n=20)
     
     ## Improve masking using CV trick
     mask = improve_mask(mask)
     
     ## Step 2 and 3: Diffusion process using mask
     output = pipe(
         prompt=qp, 
         image=init_img, 
         mask_image=Image.fromarray(mask*255).resize((512,512)), 
         generator=torch.Generator("cuda").manual_seed(100),
         num_inference_steps = steps
     ).images
     return mask , output

 p = FastDownload().download('https://raw.githubusercontent.com/johnrobinsn/diffusion_experiments/main/images/bowloberries_scaled.jpg')
 init_img = load_image(p)
 mask, output = fastDiffEdit(init_img, rp = ['Bowl of Strawberries'], qp=['Bowl of Grapes'])
 plot_diffEdit(init_img, output, mask)