import torch from libs.base_utils import do_resize_content from imagedream.ldm.util import ( instantiate_from_config, get_obj_from_str, ) from omegaconf import OmegaConf from PIL import Image import PIL import rembg class TwoStagePipeline(object): def __init__( self, stage1_model_config, stage2_model_config, stage1_sampler_config, stage2_sampler_config, device="cpu", dtype=torch.float16, resize_rate=1, ) -> None: """ only for two stage generate process. - the first stage was condition on single pixel image, gererate multi-view pixel image, based on the v2pp config - the second stage was condition on multiview pixel image generated by the first stage, generate the final image, based on the stage2-test config """ self.resize_rate = resize_rate self.stage1_model = instantiate_from_config(OmegaConf.load(stage1_model_config.config).model) self.stage1_model.load_state_dict(torch.load(stage1_model_config.resume, map_location="cpu"), strict=False) self.stage1_model = self.stage1_model.to(device).to(dtype) self.stage2_model = instantiate_from_config(OmegaConf.load(stage2_model_config.config).model) sd = torch.load(stage2_model_config.resume, map_location="cpu") self.stage2_model.load_state_dict(sd, strict=False) self.stage2_model = self.stage2_model.to(device).to(dtype) self.stage1_model.device = device self.stage2_model.device = device self.device = device self.dtype = dtype self.stage1_sampler = get_obj_from_str(stage1_sampler_config.target)( self.stage1_model, device=device, dtype=dtype, **stage1_sampler_config.params ) self.stage2_sampler = get_obj_from_str(stage2_sampler_config.target)( self.stage2_model, device=device, dtype=dtype, **stage2_sampler_config.params ) def stage1_sample( self, pixel_img, prompt="3D assets", neg_texts="uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.", step=50, scale=5, ddim_eta=0.0, ): if type(pixel_img) == str: pixel_img = Image.open(pixel_img) if isinstance(pixel_img, Image.Image): if pixel_img.mode == "RGBA": background = Image.new('RGBA', pixel_img.size, (0, 0, 0, 0)) pixel_img = Image.alpha_composite(background, pixel_img).convert("RGB") else: pixel_img = pixel_img.convert("RGB") else: raise uc = self.stage1_sampler.model.get_learned_conditioning([neg_texts]).to(self.device) stage1_images = self.stage1_sampler.i2i( self.stage1_sampler.model, self.stage1_sampler.size, prompt, uc=uc, sampler=self.stage1_sampler.sampler, ip=pixel_img, step=step, scale=scale, batch_size=self.stage1_sampler.batch_size, ddim_eta=ddim_eta, dtype=self.stage1_sampler.dtype, device=self.stage1_sampler.device, camera=self.stage1_sampler.camera, num_frames=self.stage1_sampler.num_frames, pixel_control=(self.stage1_sampler.mode == "pixel"), transform=self.stage1_sampler.image_transform, offset_noise=self.stage1_sampler.offset_noise, ) stage1_images = [Image.fromarray(img) for img in stage1_images] stage1_images.pop(self.stage1_sampler.ref_position) return stage1_images def stage2_sample(self, pixel_img, stage1_images, scale=5, step=50): if type(pixel_img) == str: pixel_img = Image.open(pixel_img) if isinstance(pixel_img, Image.Image): if pixel_img.mode == "RGBA": background = Image.new('RGBA', pixel_img.size, (0, 0, 0, 0)) pixel_img = Image.alpha_composite(background, pixel_img).convert("RGB") else: pixel_img = pixel_img.convert("RGB") else: raise stage2_images = self.stage2_sampler.i2iStage2( self.stage2_sampler.model, self.stage2_sampler.size, "3D assets", self.stage2_sampler.uc, self.stage2_sampler.sampler, pixel_images=stage1_images, ip=pixel_img, step=step, scale=scale, batch_size=self.stage2_sampler.batch_size, ddim_eta=0.0, dtype=self.stage2_sampler.dtype, device=self.stage2_sampler.device, camera=self.stage2_sampler.camera, num_frames=self.stage2_sampler.num_frames, pixel_control=(self.stage2_sampler.mode == "pixel"), transform=self.stage2_sampler.image_transform, offset_noise=self.stage2_sampler.offset_noise, ) stage2_images = [Image.fromarray(img) for img in stage2_images] return stage2_images def set_seed(self, seed): self.stage1_sampler.seed = seed self.stage2_sampler.seed = seed def __call__(self, pixel_img, prompt="3D assets", scale=5, step=50): pixel_img = do_resize_content(pixel_img, self.resize_rate) stage1_images = self.stage1_sample(pixel_img, prompt, scale=scale, step=step) stage2_images = self.stage2_sample(pixel_img, stage1_images, scale=scale, step=step) return { "ref_img": pixel_img, "stage1_images": stage1_images, "stage2_images": stage2_images, } rembg_session = rembg.new_session() def expand_to_square(image, bg_color=(0, 0, 0, 0)): # expand image to 1:1 width, height = image.size if width == height: return image new_size = (max(width, height), max(width, height)) new_image = Image.new("RGBA", new_size, bg_color) paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2) new_image.paste(image, paste_position) return new_image def remove_background( image: PIL.Image.Image, rembg_session = None, force: bool = False, **rembg_kwargs, ) -> PIL.Image.Image: do_remove = True if image.mode == "RGBA" and image.getextrema()[3][0] < 255: # explain why current do not rm bg print("alhpa channl not enpty, skip remove background, using alpha channel as mask") background = Image.new("RGBA", image.size, (0, 0, 0, 0)) image = Image.alpha_composite(background, image) do_remove = False do_remove = do_remove or force if do_remove: image = rembg.remove(image, session=rembg_session, **rembg_kwargs) return image def do_resize_content(original_image: Image, scale_rate): # resize image content wile retain the original image size if scale_rate != 1: # Calculate the new size after rescaling new_size = tuple(int(dim * scale_rate) for dim in original_image.size) # Resize the image while maintaining the aspect ratio resized_image = original_image.resize(new_size) # Create a new image with the original size and black background padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0)) paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2) padded_image.paste(resized_image, paste_position) return padded_image else: return original_image def add_background(image, bg_color=(255, 255, 255)): # given an RGBA image, alpha channel is used as mask to add background color background = Image.new("RGBA", image.size, bg_color) return Image.alpha_composite(background, image) def preprocess_image(image, background_choice, foreground_ratio, backgroud_color): """ input image is a pil image in RGBA, return RGB image """ print(background_choice) if background_choice == "Alpha as mask": background = Image.new("RGBA", image.size, (0, 0, 0, 0)) image = Image.alpha_composite(background, image) else: image = remove_background(image, rembg_session, force_remove=True) image = do_resize_content(image, foreground_ratio) image = expand_to_square(image) image = add_background(image, backgroud_color) return image.convert("RGB")