# Edit Anything trained with Stable Diffusion + ControlNet + SAM + BLIP2 from torchvision.utils import save_image from PIL import Image from pytorch_lightning import seed_everything import subprocess from collections import OrderedDict import re import cv2 import einops import gradio as gr import numpy as np import torch import random import os import requests from io import BytesIO from annotator.util import resize_image, HWC3 import torch from safetensors.torch import load_file from collections import defaultdict from diffusers import StableDiffusionControlNetPipeline from diffusers import ControlNetModel, UniPCMultistepScheduler from utils.stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline # from utils.tmp import StableDiffusionControlNetInpaintPipeline # need the latest transformers # pip install git+https://github.com/huggingface/transformers.git from transformers import AutoProcessor, Blip2ForConditionalGeneration from diffusers import ControlNetModel, DiffusionPipeline import PIL.Image # Segment-Anything init. # pip install git+https://github.com/facebookresearch/segment-anything.git try: from segment_anything import sam_model_registry, SamAutomaticMaskGenerator except ImportError: print('segment_anything not installed') result = subprocess.run( ['pip', 'install', 'git+https://github.com/facebookresearch/segment-anything.git'], check=True) print(f'Install segment_anything {result}') from segment_anything import sam_model_registry, SamAutomaticMaskGenerator if not os.path.exists('./models/sam_vit_h_4b8939.pth'): result = subprocess.run( ['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth', '-P', 'models'], check=True) print(f'Download sam_vit_h_4b8939.pth {result}') device = "cuda" if torch.cuda.is_available() else "cpu" config_dict = OrderedDict([ ('LAION Pretrained(v0-4)-SD15', 'shgao/edit-anything-v0-4-sd15'), ('LAION Pretrained(v0-4)-SD21', 'shgao/edit-anything-v0-4-sd21'), ]) def init_sam_model(): sam_checkpoint = "models/sam_vit_h_4b8939.pth" model_type = "default" sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) sam.to(device=device) sam_generator = SamAutomaticMaskGenerator(sam) return sam_generator def init_blip_processor(): blip_processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b") return blip_processor def init_blip_model(): blip_model = Blip2ForConditionalGeneration.from_pretrained( "Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto") return blip_model def get_pipeline_embeds(pipeline, prompt, negative_prompt, device): # https://github.com/huggingface/diffusers/issues/2136 """ Get pipeline embeds for prompts bigger than the maxlength of the pipe :param pipeline: :param prompt: :param negative_prompt: :param device: :return: """ max_length = pipeline.tokenizer.model_max_length # simple way to determine length of tokens count_prompt = len(re.split(r', ', prompt)) count_negative_prompt = len(re.split(r', ', negative_prompt)) # create the tensor based on which prompt is longer if count_prompt >= count_negative_prompt: input_ids = pipeline.tokenizer( prompt, return_tensors="pt", truncation=False).input_ids.to(device) shape_max_length = input_ids.shape[-1] negative_ids = pipeline.tokenizer(negative_prompt, truncation=False, padding="max_length", max_length=shape_max_length, return_tensors="pt").input_ids.to(device) else: negative_ids = pipeline.tokenizer( negative_prompt, return_tensors="pt", truncation=False).input_ids.to(device) shape_max_length = negative_ids.shape[-1] input_ids = pipeline.tokenizer(prompt, return_tensors="pt", truncation=False, padding="max_length", max_length=shape_max_length).input_ids.to(device) concat_embeds = [] neg_embeds = [] for i in range(0, shape_max_length, max_length): concat_embeds.append(pipeline.text_encoder( input_ids[:, i: i + max_length])[0]) neg_embeds.append(pipeline.text_encoder( negative_ids[:, i: i + max_length])[0]) return torch.cat(concat_embeds, dim=1), torch.cat(neg_embeds, dim=1) def load_lora_weights(pipeline, checkpoint_path, multiplier, device, dtype): LORA_PREFIX_UNET = "lora_unet" LORA_PREFIX_TEXT_ENCODER = "lora_te" # load LoRA weight from .safetensors if isinstance(checkpoint_path, str): state_dict = load_file(checkpoint_path, device=device) updates = defaultdict(dict) for key, value in state_dict.items(): # it is suggested to print out the key, it usually will be something like below # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" layer, elem = key.split('.', 1) updates[layer][elem] = value # directly update weight in diffusers model for layer, elems in updates.items(): if "text" in layer: layer_infos = layer.split( LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") curr_layer = pipeline.text_encoder else: layer_infos = layer.split( LORA_PREFIX_UNET + "_")[-1].split("_") curr_layer = pipeline.unet # find the target layer temp_name = layer_infos.pop(0) while len(layer_infos) > -1: try: curr_layer = curr_layer.__getattr__(temp_name) if len(layer_infos) > 0: temp_name = layer_infos.pop(0) elif len(layer_infos) == 0: break except Exception: if len(temp_name) > 0: temp_name += "_" + layer_infos.pop(0) else: temp_name = layer_infos.pop(0) # get elements for this layer weight_up = elems['lora_up.weight'].to(dtype) weight_down = elems['lora_down.weight'].to(dtype) alpha = elems['alpha'] if alpha: alpha = alpha.item() / weight_up.shape[1] else: alpha = 1.0 # update weight if len(weight_up.shape) == 4: curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up.squeeze( 3).squeeze(2), weight_down.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) else: curr_layer.weight.data += multiplier * \ alpha * torch.mm(weight_up, weight_down) else: for ckptpath in checkpoint_path: state_dict = load_file(ckptpath, device=device) updates = defaultdict(dict) for key, value in state_dict.items(): # it is suggested to print out the key, it usually will be something like below # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" layer, elem = key.split('.', 1) updates[layer][elem] = value # directly update weight in diffusers model for layer, elems in updates.items(): if "text" in layer: layer_infos = layer.split( LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") curr_layer = pipeline.text_encoder else: layer_infos = layer.split( LORA_PREFIX_UNET + "_")[-1].split("_") curr_layer = pipeline.unet # find the target layer temp_name = layer_infos.pop(0) while len(layer_infos) > -1: try: curr_layer = curr_layer.__getattr__(temp_name) if len(layer_infos) > 0: temp_name = layer_infos.pop(0) elif len(layer_infos) == 0: break except Exception: if len(temp_name) > 0: temp_name += "_" + layer_infos.pop(0) else: temp_name = layer_infos.pop(0) # get elements for this layer weight_up = elems['lora_up.weight'].to(dtype) weight_down = elems['lora_down.weight'].to(dtype) alpha = elems['alpha'] if alpha: alpha = alpha.item() / weight_up.shape[1] else: alpha = 1.0 # update weight if len(weight_up.shape) == 4: curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up.squeeze( 3).squeeze(2), weight_down.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3) else: curr_layer.weight.data += multiplier * \ alpha * torch.mm(weight_up, weight_down) return pipeline def make_inpaint_condition(image, image_mask): # image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 image = image / 255.0 print("img", image.max(), image.min(), image_mask.max(), image_mask.min()) # image_mask = np.array(image_mask.convert("L")) assert image.shape[0:1] == image_mask.shape[0: 1], "image and image_mask must have the same image size" image[image_mask > 128] = -1.0 # set as masked pixel image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) image = torch.from_numpy(image) return image def obtain_generation_model(base_model_path, lora_model_path, controlnet_path, generation_only=False, extra_inpaint=True, lora_weight=1.0): controlnet = [] controlnet.append(ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)) # sam control if (not generation_only) and extra_inpaint: # inpainting control print("Warning: ControlNet based inpainting model only support SD1.5 for now.") controlnet.append( ControlNetModel.from_pretrained( 'lllyasviel/control_v11p_sd15_inpaint', torch_dtype=torch.float16) # inpainting controlnet ) if generation_only: pipe = StableDiffusionControlNetPipeline.from_pretrained( base_model_path, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None ) else: pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( base_model_path, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None ) if lora_model_path is not None: pipe = load_lora_weights( pipe, [lora_model_path], lora_weight, 'cpu', torch.float32) # speed up diffusion process with faster scheduler and memory optimization pipe.scheduler = UniPCMultistepScheduler.from_config( pipe.scheduler.config) # remove following line if xformers is not installed pipe.enable_xformers_memory_efficient_attention() pipe.enable_model_cpu_offload() return pipe def obtain_tile_model(base_model_path, lora_model_path, lora_weight=1.0): controlnet = ControlNetModel.from_pretrained( 'lllyasviel/control_v11f1e_sd15_tile', torch_dtype=torch.float16) # tile controlnet if base_model_path=='runwayml/stable-diffusion-v1-5' or base_model_path=='stabilityai/stable-diffusion-2-inpainting': print("base_model_path", base_model_path) pipe = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None ) else: pipe = StableDiffusionControlNetPipeline.from_pretrained( base_model_path, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None ) if lora_model_path is not None: pipe = load_lora_weights( pipe, [lora_model_path], lora_weight, 'cpu', torch.float32) # speed up diffusion process with faster scheduler and memory optimization pipe.scheduler = UniPCMultistepScheduler.from_config( pipe.scheduler.config) # remove following line if xformers is not installed pipe.enable_xformers_memory_efficient_attention() pipe.enable_model_cpu_offload() return pipe def show_anns(anns): if len(anns) == 0: return sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True) full_img = None # for ann in sorted_anns: for i in range(len(sorted_anns)): ann = anns[i] m = ann['segmentation'] if full_img is None: full_img = np.zeros((m.shape[0], m.shape[1], 3)) map = np.zeros((m.shape[0], m.shape[1]), dtype=np.uint16) map[m != 0] = i + 1 color_mask = np.random.random((1, 3)).tolist()[0] full_img[m != 0] = color_mask full_img = full_img*255 # anno encoding from https://github.com/LUSSeg/ImageNet-S res = np.zeros((map.shape[0], map.shape[1], 3)) res[:, :, 0] = map % 256 res[:, :, 1] = map // 256 res.astype(np.float32) full_img = Image.fromarray(np.uint8(full_img)) return full_img, res class EditAnythingLoraModel: def __init__(self, base_model_path='../chilloutmix_NiPrunedFp32Fix', lora_model_path='../40806/mix4', use_blip=True, blip_processor=None, blip_model=None, sam_generator=None, controlmodel_name='LAION Pretrained(v0-4)-SD15', extra_inpaint=True, # used when the base model is not an inpainting model. tile_model=None, lora_weight=1.0, ): self.device = device self.use_blip = use_blip # Diffusion init using diffusers. self.default_controlnet_path = config_dict[controlmodel_name] self.base_model_path = base_model_path self.lora_model_path = lora_model_path self.defalut_enable_all_generate = False self.extra_inpaint = extra_inpaint self.pipe = obtain_generation_model( base_model_path, lora_model_path, self.default_controlnet_path, generation_only=False, extra_inpaint=extra_inpaint, lora_weight=lora_weight) # Segment-Anything init. if sam_generator is not None: self.sam_generator = sam_generator else: self.sam_generator = init_sam_model() # BLIP2 init. if use_blip: if blip_processor is not None: self.blip_processor = blip_processor else: self.blip_processor = init_blip_processor() if blip_model is not None: self.blip_model = blip_model else: self.blip_model = init_blip_model() # tile model init. if tile_model is not None: self.tile_pipe = tile_model else: self.tile_pipe = obtain_tile_model(base_model_path, lora_model_path, lora_weight=lora_weight) def get_blip2_text(self, image): inputs = self.blip_processor(image, return_tensors="pt").to( self.device, torch.float16) generated_ids = self.blip_model.generate(**inputs, max_new_tokens=50) generated_text = self.blip_processor.batch_decode( generated_ids, skip_special_tokens=True)[0].strip() return generated_text def get_sam_control(self, image): masks = self.sam_generator.generate(image) full_img, res = show_anns(masks) return full_img, res @torch.inference_mode() def process(self, source_image, enable_all_generate, mask_image, control_scale, enable_auto_prompt, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, enable_tile=True, condition_model=None): if condition_model is None: this_controlnet_path = self.default_controlnet_path else: this_controlnet_path = config_dict[condition_model] input_image = source_image["image"] if mask_image is None: if enable_all_generate != self.defalut_enable_all_generate: self.pipe = obtain_generation_model( self.base_model_path, self.lora_model_path, this_controlnet_path, enable_all_generate, self.extra_inpaint) self.defalut_enable_all_generate = enable_all_generate if enable_all_generate: print("source_image", source_image["mask"].shape, input_image.shape,) mask_image = np.ones( (input_image.shape[0], input_image.shape[1], 3))*255 else: mask_image = source_image["mask"] if self.default_controlnet_path != this_controlnet_path: print("To Use:", this_controlnet_path, "Current:", self.default_controlnet_path) print("Change condition model to:", this_controlnet_path) self.pipe = obtain_generation_model( self.base_model_path, self.lora_model_path, this_controlnet_path, enable_all_generate, self.extra_inpaint) self.default_controlnet_path = this_controlnet_path torch.cuda.empty_cache() with torch.no_grad(): if self.use_blip and enable_auto_prompt: print("Generating text:") blip2_prompt = self.get_blip2_text(input_image) print("Generated text:", blip2_prompt) if len(prompt) > 0: prompt = blip2_prompt + ',' + prompt else: prompt = blip2_prompt input_image = HWC3(input_image) img = resize_image(input_image, image_resolution) H, W, C = img.shape print("Generating SAM seg:") # the default SAM model is trained with 1024 size. full_segmask, detected_map = self.get_sam_control( resize_image(input_image, detect_resolution)) detected_map = HWC3(detected_map.astype(np.uint8)) detected_map = cv2.resize( detected_map, (W, H), interpolation=cv2.INTER_LINEAR) control = torch.from_numpy( detected_map.copy()).float().cuda() control = torch.stack([control for _ in range(num_samples)], dim=0) control = einops.rearrange(control, 'b h w c -> b c h w').clone() mask_image = HWC3(mask_image.astype(np.uint8)) mask_image_tmp = cv2.resize( mask_image, (W, H), interpolation=cv2.INTER_LINEAR) mask_image = Image.fromarray(mask_image_tmp) if seed == -1: seed = random.randint(0, 65535) seed_everything(seed) generator = torch.manual_seed(seed) postive_prompt = prompt + ', ' + a_prompt negative_prompt = n_prompt prompt_embeds, negative_prompt_embeds = get_pipeline_embeds( self.pipe, postive_prompt, negative_prompt, "cuda") prompt_embeds = torch.cat([prompt_embeds] * num_samples, dim=0) negative_prompt_embeds = torch.cat( [negative_prompt_embeds] * num_samples, dim=0) if enable_all_generate and self.extra_inpaint: self.pipe.safety_checker = lambda images, clip_input: ( images, False) x_samples = self.pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_images_per_prompt=num_samples, num_inference_steps=ddim_steps, generator=generator, height=H, width=W, image=[control.type(torch.float16)], controlnet_conditioning_scale=[float(control_scale)], ).images else: multi_condition_image = [] multi_condition_scale = [] multi_condition_image.append(control.type(torch.float16)) multi_condition_scale.append(float(control_scale)) if self.extra_inpaint: inpaint_image = make_inpaint_condition(img, mask_image_tmp) print(inpaint_image.shape) multi_condition_image.append(inpaint_image.type(torch.float16)) multi_condition_scale.append(1.0) x_samples = self.pipe( image=img, mask_image=mask_image, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_images_per_prompt=num_samples, num_inference_steps=ddim_steps, generator=generator, controlnet_conditioning_image=multi_condition_image, height=H, width=W, controlnet_conditioning_scale=multi_condition_scale, ).images results = [x_samples[i] for i in range(num_samples)] if True: img_tile = [PIL.Image.fromarray(resize_image(np.array(x_samples[i]), 1024)) for i in range(num_samples)] # for each in img_tile: # print("tile",each.size) prompt_embeds, negative_prompt_embeds = get_pipeline_embeds( self.tile_pipe, postive_prompt, negative_prompt, "cuda") prompt_embeds = torch.cat([prompt_embeds] * num_samples, dim=0) negative_prompt_embeds = torch.cat( [negative_prompt_embeds] * num_samples, dim=0) x_samples_tile = self.tile_pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, num_images_per_prompt=num_samples, num_inference_steps=ddim_steps, generator=generator, height=img_tile[0].size[1], width=img_tile[0].size[0], image=img_tile, controlnet_conditioning_scale=1.0, ).images results_tile = [x_samples_tile[i] for i in range(num_samples)] results = results_tile + results return [full_segmask, mask_image] + results, prompt def download_image(url): response = requests.get(url) return Image.open(BytesIO(response.content)).convert("RGB")