diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index c7a0c83a..6cdfaeaa 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -82,6 +82,7 @@ class Models(object): image_skychange = 'image-skychange' video_human_matting = 'video-human-matting' human_reconstruction = 'human-reconstruction' + text_texture_generation = 'text-texture-generation' video_frame_interpolation = 'video-frame-interpolation' video_object_segmentation = 'video-object-segmentation' video_deinterlace = 'video-deinterlace' @@ -406,6 +407,7 @@ class Pipelines(object): image_skychange = 'image-skychange' video_human_matting = 'video-human-matting' human_reconstruction = 'human-reconstruction' + text_texture_generation = 'text-texture-generation' vision_middleware_multi_task = 'vision-middleware-multi-task' vidt = 'vidt' video_frame_interpolation = 'video-frame-interpolation' @@ -839,6 +841,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { 'damo/cv_effnetv2_video-human-matting'), Tasks.human_reconstruction: (Pipelines.human_reconstruction, 'damo/cv_hrnet_image-human-reconstruction'), + Tasks.text_texture_generation: ( + Pipelines.text_texture_generation, + 'damo/cv_diffuser_text-texture-generation'), Tasks.video_frame_interpolation: ( Pipelines.video_frame_interpolation, 'damo/cv_raft_video-frame-interpolation'), diff --git a/modelscope/models/cv/__init__.py b/modelscope/models/cv/__init__.py index 2417f33b..5cbee709 100644 --- a/modelscope/models/cv/__init__.py +++ b/modelscope/models/cv/__init__.py @@ -23,7 +23,8 @@ from . import (action_recognition, animal_recognition, bad_image_detecting, referring_video_object_segmentation, robust_image_classification, salient_detection, shop_segmentation, stream_yolo, super_resolution, - surface_recon_common, table_recognition, video_deinterlace, + surface_recon_common, table_recognition, + text_texture_generation, video_deinterlace, video_frame_interpolation, video_object_segmentation, video_panoptic_segmentation, video_single_object_tracking, video_stabilization, video_summarization, diff --git a/modelscope/models/cv/text_texture_generation/Tex2Texture.py b/modelscope/models/cv/text_texture_generation/Tex2Texture.py new file mode 100644 index 00000000..5e6eee3e --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/Tex2Texture.py @@ -0,0 +1,660 @@ +# Copyright © Alibaba, Inc. and its affiliates. +# The implementation here is modifed based on StableDiffusionControlNetInpaintPipeline, +# originally Apache 2.0 License and public available at +# https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py + +import os +from typing import Any, Callable, Dict, List, Optional, Union + +import cv2 +import numpy as np +import PIL +import PIL.Image as Image +import torch +import torchvision.transforms as transforms +from diffusers import (AutoencoderKL, ControlNetModel, DiffusionPipeline, + EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, + StableDiffusionControlNetImg2ImgPipeline, + StableDiffusionControlNetInpaintPipeline, + StableDiffusionInpaintPipeline, StableDiffusionPipeline, + UNet2DConditionModel) +from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel +from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.utils import (deprecate, is_accelerate_available, + is_accelerate_version, is_compiled_module, + logging, randn_tensor, replace_example_docstring) +from pytorch3d.io import load_obj, load_objs_as_meshes, save_obj + +from modelscope.metainfo import Models +from modelscope.models.base import Tensor, TorchModel +from modelscope.models.builder import MODELS +from modelscope.models.cv.text_texture_generation.lib2.camera import * +from modelscope.models.cv.text_texture_generation.lib2.init_view import * +from modelscope.models.cv.text_texture_generation.utils import * +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler + >>> from diffusers.utils import load_image + >>> import numpy as np + >>> import torch + + >>> init_image = load_image(image_path) + >>> init_image = init_image.resize((512, 512)) + >>> generator = torch.Generator(device="cpu").manual_seed(1) + >>> mask_image = load_image(mask_path) + >>> mask_image = mask_image.resize((512, 512)) + >>> def make_inpaint_condition(image, image_mask): + ... image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 + ... image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0 + ... assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size" + ... image[image_mask > 0.5] = -1.0 # set as masked pixel + ... image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) + ... image = torch.from_numpy(image) + ... return image + >>> control_image = make_inpaint_condition(init_image, mask_image) + >>> controlnet = ControlNetModel.from_pretrained( + ... "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 + ... ) + >>> pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained( + ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 + ... ) + >>> pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) + >>> pipe.enable_model_cpu_offload() + >>> image = pipe( + ... "a handsome man with ray-ban sunglasses", + ... num_inference_steps=20, + ... generator=generator, + ... eta=1.0, + ... image=init_image, + ... mask_image=mask_image, + ... control_image=control_image, + ... ).images[0] + ``` +""" + + +@MODELS.register_module( + Tasks.text_texture_generation, module_name=Models.text_texture_generation) +class Tex2Texture(TorchModel): + + def __init__(self, model_dir, *args, **kwargs): + """The Tex2Texture is modified based on TEXTure and Text2Tex, publicly available at + https://github.com/TEXTurePaper/TEXTurePaper & + https://github.com/daveredrum/Text2Tex + Args: + model_dir: the root directory of the model files + """ + super().__init__(model_dir=model_dir, *args, **kwargs) + if torch.cuda.is_available(): + self.device = torch.device('cuda') + logger.info('Use GPU: {}'.format(self.device)) + else: + print('no gpu avaiable') + exit() + + model_path = model_dir + '/base_model/' + controlmodel_path = model_dir + '/control_model/' + inpaintmodel_path = model_dir + '/inpaint_model/' + torch_dtype = kwargs.get('torch_dtype', torch.float16) + self.controlnet = ControlNetModel.from_pretrained( + controlmodel_path, torch_dtype=torch_dtype).to(self.device) + self.inpaintmodel = StableDiffusionInpaintPipeline.from_pretrained( + inpaintmodel_path, + torch_dtype=torch_dtype, + ).to(self.device) + self.pipe = StableDiffusionControlinpaintPipeline.from_pretrained( + model_path, controlnet=self.controlnet, + torch_dtype=torch_dtype).to(self.device) + logger.info('model load over') + + def init_mesh(self, mesh_path): + verts, faces, aux = load_obj(mesh_path, device=self.device) + mesh = load_objs_as_meshes([mesh_path], device=self.device) + return mesh, verts, faces, aux + + def normalize_mesh(self, mesh): + bbox = mesh.get_bounding_boxes() + num_verts = mesh.verts_packed().shape[0] + mesh_center = bbox.mean(dim=2).repeat(num_verts, 1) + mesh = mesh.offset_verts(-mesh_center) + lens = bbox[0, :, 1] - bbox[0, :, 0] + max_len = lens.max() + scale = 0.9 / max_len + scale = scale.unsqueeze(0).repeat(num_verts) + # mesh.scale_verts_(scale) + new_mesh = mesh.scale_verts(scale) + return new_mesh.verts_packed(), new_mesh, mesh_center, scale + + def save_normalized_obj(self, verts, faces, aux, path='normalized.obj'): + print('=> saving normalized mesh file...') + obj_path = path + save_obj( + obj_path, + verts=verts, + faces=faces.verts_idx, + decimal_places=5, + verts_uvs=aux.verts_uvs, + faces_uvs=faces.textures_idx, + texture_map=aux.texture_images[list(aux.texture_images.keys())[0]]) + + def mesh_normalized(self, mesh_path, save_path='normalized.obj'): + mesh, verts, faces, aux = self.init_mesh(mesh_path) + verts, mesh, mesh_center, scale = self.normalize_mesh(mesh) + self.save_normalized_obj(verts, faces, aux, save_path) + return mesh, verts, faces, aux, mesh_center, scale + + +def prepare_mask_and_masked_image(image, + mask, + height, + width, + return_image=False): + if image is None: + raise ValueError('`image` input cannot be undefined.') + + if mask is None: + raise ValueError('`mask_image` input cannot be undefined.') + + if isinstance(image, torch.Tensor): + if not isinstance(mask, torch.Tensor): + raise TypeError( + f'`image` is a torch.Tensor but `mask` (type: {type(mask)} is not' + ) + + # Batch single image + if image.ndim == 3: + assert image.shape[ + 0] == 3, 'Image outside a batch should be of shape (3, H, W)' + image = image.unsqueeze(0) + + # Batch and add channel dim for single mask + if mask.ndim == 2: + mask = mask.unsqueeze(0).unsqueeze(0) + + # Batch single mask or add channel dim + if mask.ndim == 3: + # Single batched mask, no channel dim or single mask not batched but channel dim + if mask.shape[0] == 1: + mask = mask.unsqueeze(0) + + # Batched masks no channel dim + else: + mask = mask.unsqueeze(1) + + assert image.ndim == 4 and mask.ndim == 4, 'Image and Mask must have 4 dimensions' + assert image.shape[-2:] == mask.shape[ + -2:], 'Image and Mask must have the same spatial dimensions' + assert image.shape[0] == mask.shape[ + 0], 'Image and Mask must have the same batch size' + + # Check image is in [-1, 1] + if image.min() < -1 or image.max() > 1: + raise ValueError('Image should be in [-1, 1] range') + + # Check mask is in [0, 1] + if mask.min() < 0 or mask.max() > 1: + raise ValueError('Mask should be in [0, 1] range') + + # Binarize mask + mask[mask < 0.5] = 0 + mask[mask >= 0.5] = 1 + + # Image as float32 + image = image.to(dtype=torch.float32) + elif isinstance(mask, torch.Tensor): + raise TypeError( + f'`mask` is a torch.Tensor but `image` (type: {type(image)} is not' + ) + else: + # preprocess image + if isinstance(image, (PIL.Image.Image, np.ndarray)): + image = [image] + if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): + # resize all images w.r.t passed height an width + image = [ + i.resize((width, height), resample=PIL.Image.LANCZOS) + for i in image + ] + image = [np.array(i.convert('RGB'))[None, :] for i in image] + image = np.concatenate(image, axis=0) + elif isinstance(image, list) and isinstance(image[0], np.ndarray): + image = np.concatenate([i[None, :] for i in image], axis=0) + + image = image.transpose(0, 3, 1, 2) + image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0 + + # preprocess mask + if isinstance(mask, (PIL.Image.Image, np.ndarray)): + mask = [mask] + + if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image): + mask = [ + i.resize((width, height), resample=PIL.Image.LANCZOS) + for i in mask + ] + mask = np.concatenate( + [np.array(m.convert('L'))[None, None, :] for m in mask], + axis=0) + mask = mask.astype(np.float32) / 255.0 + elif isinstance(mask, list) and isinstance(mask[0], np.ndarray): + mask = np.concatenate([m[None, None, :] for m in mask], axis=0) + + mask[mask < 0.5] = 0 + mask[mask >= 0.5] = 1 + mask = torch.from_numpy(mask) + + masked_image = image * (mask < 0.5) + + # n.b. ensure backwards compatibility as old function does not return image + if return_image: + return mask, masked_image, image + + return mask, masked_image + + +class StableDiffusionControlinpaintPipeline( + StableDiffusionControlNetInpaintPipeline): + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: Union[torch.Tensor, PIL.Image.Image] = None, + mask_image: Union[torch.Tensor, PIL.Image.Image] = None, + control_image: Union[torch.FloatTensor, PIL.Image.Image, np.ndarray, + List[torch.FloatTensor], List[PIL.Image.Image], + List[np.ndarray], ] = None, + height: Optional[int] = None, + width: Optional[int] = None, + strength: float = 1.0, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, + List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = 'pil', + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], + None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_conditioning_scale: Union[float, List[float]] = 0.5, + guess_mode: bool = False, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. + instead. + image (`torch.FloatTensor`, `PIL.Image.Image`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, + `List[List[torch.FloatTensor]]`, or `List[List[PIL.Image.Image]]`): + The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If + the type is specified as `Torch.FloatTensor`, it is passed to ControlNet as is. `PIL.Image.Image` can + also be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If + height and/or width are passed, `image` is resized according to them. If multiple ControlNets are + specified in init, images must be passed as a list such that each element of the list can be correctly + batched for input to a single controlnet. + height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. + strength (`float`, *optional*, defaults to 1.): + Conceptually, indicates how much to transform the masked portion of the reference `image`. Must be + between 0 and 1. `image` will be used as a starting point, adding more noise to it the larger the + `strength`. The number of denoising steps depends on the amount of noise initially added. When + `strength` is 1, added noise will be maximum and the denoising process will run for the full number of + iterations specified in `num_inference_steps`. A value of 1, therefore, essentially ignores the masked + portion of the reference `image`. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + guidance_scale (`float`, *optional*, defaults to 7.5): + Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). + `guidance_scale` is defined as `w` of equation 2. of [Imagen + Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > + 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, + usually at the expense of lower image quality. + negative_prompt (`str` or `List[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to + [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) + to make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt + weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input + argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between + [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a + plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be + called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). + controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 0.5): + The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added + to the residual in the original unet. If multiple ControlNets are specified in init, you can set the + corresponding scale as a list. Note that by default, we use a smaller conditioning scale for inpainting + than for [`~StableDiffusionControlNetPipeline.__call__`]. + guess_mode (`bool`, *optional*, defaults to `False`): + In this mode, the ControlNet encoder will try best to recognize the content of the input image even if + you remove all prompts. The `guidance_scale` between 3.0 and 5.0 is recommended. + + Examples: + + Returns: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: + [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. + When returning a tuple, the first element is a list with the generated images, and the second element is a + list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" + (nsfw) content, according to the `safety_checker`. + """ + # 0. Default height and width to unet + height, width = self._default_height_width(height, width, image) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + control_image, + height, + width, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + controlnet_conditioning_scale, + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + controlnet = self.controlnet._orig_mod if is_compiled_module( + self.controlnet) else self.controlnet + + if isinstance(controlnet, MultiControlNetModel) and isinstance( + controlnet_conditioning_scale, float): + controlnet_conditioning_scale = [controlnet_conditioning_scale + ] * len(controlnet.nets) + + global_pool_conditions = ( + controlnet.config.global_pool_conditions if isinstance( + controlnet, ControlNetModel) else + controlnet.nets[0].config.global_pool_conditions) + guess_mode = guess_mode or global_pool_conditions + + # 3. Encode input prompt + text_encoder_lora_scale = ( + cross_attention_kwargs.get('scale', None) + if cross_attention_kwargs is not None else None) + prompt_embeds = self._encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + + # 4. Prepare image + if isinstance(controlnet, ControlNetModel): + control_image = self.prepare_control_image( + image=control_image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + guess_mode=guess_mode, + ) + elif isinstance(controlnet, MultiControlNetModel): + control_images = [] + + for control_image_ in control_image: + control_image_ = self.prepare_control_image( + image=control_image_, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + guess_mode=guess_mode, + ) + + control_images.append(control_image_) + + control_image = control_images + else: + assert False + + # 4. Preprocess mask and image - resizes image and mask w.r.t height and width + mask, masked_image, init_image = prepare_mask_and_masked_image( + image, mask_image, height, width, return_image=True) + + # 5. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps, num_inference_steps = self.get_timesteps( + num_inference_steps=num_inference_steps, + strength=strength, + device=device) + # at which timestep to set the initial noise (n.b. 50% if strength is 0.5) + latent_timestep = timesteps[:1].repeat(batch_size + * num_images_per_prompt) + # create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise + is_strength_max = strength == 1.0 + + # 6. Prepare latent variables + num_channels_latents = self.vae.config.latent_channels + num_channels_unet = self.unet.config.in_channels + return_image_latents = num_channels_unet == 4 + latents_outputs = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + image=init_image, + timestep=latent_timestep, + is_strength_max=is_strength_max, + return_noise=True, + return_image_latents=return_image_latents, + ) + + if return_image_latents: + latents, noise, image_latents = latents_outputs + else: + latents, noise = latents_outputs + + # 7. Prepare mask latent variables + mask, masked_image_latents = self.prepare_mask_latents( + mask, + masked_image, + batch_size * num_images_per_prompt, + height, + width, + prompt_embeds.dtype, + device, + generator, + do_classifier_free_guidance, + ) + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 8. Denoising loop + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat( + [latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t) + + # controlnet(s) inference + if guess_mode and do_classifier_free_guidance: + # Infer ControlNet only for the conditional batch. + control_model_input = latents + control_model_input = self.scheduler.scale_model_input( + control_model_input, t) + controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] + else: + control_model_input = latent_model_input + controlnet_prompt_embeds = prompt_embeds + + down_block_res_samples, mid_block_res_sample = self.controlnet( + control_model_input, + t, + encoder_hidden_states=controlnet_prompt_embeds, + controlnet_cond=control_image, + conditioning_scale=controlnet_conditioning_scale, + guess_mode=guess_mode, + return_dict=False, + ) + + if guess_mode and do_classifier_free_guidance: + # Infered ControlNet only for the conditional batch. + # To apply the output of ControlNet to both the unconditional and conditional batches, + # add 0 to the unconditional batch to keep it unchanged. + down_block_res_samples = [ + torch.cat([torch.zeros_like(d), d]) + for d in down_block_res_samples + ] + mid_block_res_sample = torch.cat([ + torch.zeros_like(mid_block_res_sample), + mid_block_res_sample + ]) + + # predict the noise residual + if num_channels_unet == 9: + latent_model_input = torch.cat( + [latent_model_input, mask, masked_image_latents], + dim=1) + + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step( + noise_pred, + t, + latents, + **extra_step_kwargs, + return_dict=False)[0] + + if num_channels_unet == 4: + init_latents_proper = image_latents[:1] + init_mask = mask[:1] + + if i < len(timesteps) - 1: + init_latents_proper = self.scheduler.add_noise( + init_latents_proper, noise, torch.tensor([t])) + + latents = (1 - init_mask + ) * init_latents_proper + init_mask * latents + + if i == len(timesteps) - 1 or ((i + 1) % self.scheduler.order + == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + # If we do sequential model offloading, let's offload unet and controlnet + # manually for max memory savings + if hasattr( + self, + 'final_offload_hook') and self.final_offload_hook is not None: + self.unet.to('cpu') + self.controlnet.to('cpu') + torch.cuda.empty_cache() + + if not output_type == 'latent': + image = self.vae.decode( + latents / self.vae.config.scaling_factor, return_dict=False)[0] + image, has_nsfw_concept = self.run_safety_checker( + image, device, prompt_embeds.dtype) + else: + image = latents + has_nsfw_concept = None + + if has_nsfw_concept is None: + do_denormalize = [True] * image.shape[0] + else: + do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + + image = self.image_processor.postprocess( + image, output_type=output_type, do_denormalize=do_denormalize) + + if hasattr( + self, + 'final_offload_hook') and self.final_offload_hook is not None: + self.final_offload_hook.offload() + if not return_dict: + return (image, has_nsfw_concept) + return StableDiffusionPipelineOutput( + images=image, nsfw_content_detected=has_nsfw_concept) diff --git a/modelscope/models/cv/text_texture_generation/__init__.py b/modelscope/models/cv/text_texture_generation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/text_texture_generation/lib2/__init__.py b/modelscope/models/cv/text_texture_generation/lib2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/text_texture_generation/lib2/camera.py b/modelscope/models/cv/text_texture_generation/lib2/camera.py new file mode 100644 index 00000000..cd36bcd9 --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/lib2/camera.py @@ -0,0 +1,165 @@ +# customized +import sys + +import numpy as np +import torch +from pytorch3d.renderer import PerspectiveCameras, look_at_view_transform +from sklearn.metrics.pairwise import cosine_similarity + +from modelscope.models.cv.text_texture_generation.lib2.init_view import \ + VIEWPOINTS + +sys.path.append('.') + +# ---------------- UTILS ---------------------- + + +def degree_to_radian(d): + return d * np.pi / 180 + + +def radian_to_degree(r): + return 180 * r / np.pi + + +def xyz_to_polar(xyz): + """ assume y-axis is the up axis """ + + x, y, z = xyz + + theta = 180 * np.arccos(z) / np.pi + phi = 180 * np.arccos(y) / np.pi + + return theta, phi + + +def polar_to_xyz(theta, phi, dist): + """ assume y-axis is the up axis """ + + theta = degree_to_radian(theta) + phi = degree_to_radian(phi) + + x = np.sin(phi) * np.sin(theta) * dist + y = np.cos(phi) * dist + z = np.sin(phi) * np.cos(theta) * dist + + return [x, y, z] + + +# ---------------- VIEWPOINTS ---------------------- + + +def filter_viewpoints(pre_viewpoints: dict, viewpoints: dict): + """ return the binary mask of viewpoints to be filtered """ + + filter_mask = [0 for _ in viewpoints.keys()] + for i, v in viewpoints.items(): + x_v, y_v, z_v = polar_to_xyz(v['azim'], 90 - v['elev'], v['dist']) + + for _, pv in pre_viewpoints.items(): + x_pv, y_pv, z_pv = polar_to_xyz(pv['azim'], 90 - pv['elev'], + pv['dist']) + sim = cosine_similarity( + np.array([[x_v, y_v, z_v]]), np.array([[x_pv, y_pv, z_pv]]))[0, + 0] + + if sim > 0.9: + filter_mask[i] = 1 + + return filter_mask + + +def init_viewpoints(init_dist, + init_elev, + init_azim, + use_principle=True, + use_shapenet=False, + use_objaverse=False): + sample_space = 12 + (dist_list, elev_list, azim_list, + sector_list) = init_predefined_viewpoints(sample_space, init_dist, + init_elev) + + # punishments for views -> in case always selecting the same view + view_punishments = [1 for _ in range(len(dist_list))] + + if use_principle: + (dist_list, elev_list, azim_list, sector_list, + view_punishments) = init_principle_viewpoints(dist_list, elev_list, + azim_list, sector_list, + view_punishments, + use_shapenet, + use_objaverse) + azim_list = [v - init_azim for v in azim_list] + elev_list = [v - init_elev for v in elev_list] + + return dist_list, elev_list, azim_list, sector_list, view_punishments + + +def init_principle_viewpoints(dist_list, + elev_list, + azim_list, + sector_list, + view_punishments, + use_shapenet=False, + use_objaverse=False): + if use_shapenet: + key = 'shapenet' + + pre_elev_list = [v for v in VIEWPOINTS[key]['elev']] + pre_azim_list = [v for v in VIEWPOINTS[key]['azim']] + pre_sector_list = [v for v in VIEWPOINTS[key]['sector']] + + num_principle = 10 + pre_dist_list = [dist_list[0] for _ in range(num_principle)] + pre_view_punishments = [0 for _ in range(num_principle)] + + elif use_objaverse: + key = 'objaverse' + + pre_elev_list = [v for v in VIEWPOINTS[key]['elev']] + pre_azim_list = [v for v in VIEWPOINTS[key]['azim']] + pre_sector_list = [v for v in VIEWPOINTS[key]['sector']] + + num_principle = 10 + pre_dist_list = [dist_list[0] for _ in range(num_principle)] + pre_view_punishments = [0 for _ in range(num_principle)] + else: + num_principle = 12 + pre_elev_list = [v for v in VIEWPOINTS[num_principle]['elev']] + pre_azim_list = [v for v in VIEWPOINTS[num_principle]['azim']] + pre_sector_list = [v for v in VIEWPOINTS[num_principle]['sector']] + pre_dist_list = [dist_list[0] for _ in range(num_principle)] + pre_view_punishments = [0 for _ in range(num_principle)] + + dist_list = pre_dist_list + dist_list + elev_list = pre_elev_list + elev_list + azim_list = pre_azim_list + azim_list + sector_list = pre_sector_list + sector_list + view_punishments = pre_view_punishments + view_punishments + + return dist_list, elev_list, azim_list, sector_list, view_punishments + + +def init_predefined_viewpoints(sample_space, init_dist, init_elev): + viewpoints = VIEWPOINTS[sample_space] + + assert sample_space == len(viewpoints['sector']) + + dist_list = [init_dist + for _ in range(sample_space)] # always the same dist + elev_list = [viewpoints['elev'][i] for i in range(sample_space)] + azim_list = [viewpoints['azim'][i] for i in range(sample_space)] + sector_list = [viewpoints['sector'][i] for i in range(sample_space)] + + return dist_list, elev_list, azim_list, sector_list + + +def init_camera(dist, elev, azim, image_size, device): + R, T = look_at_view_transform(dist, elev, azim) + image_size = torch.tensor([image_size, image_size]).unsqueeze(0) + T[0][2] = dist + cameras = PerspectiveCameras( + R=R, T=T, device=device, image_size=image_size) + + return cameras diff --git a/modelscope/models/cv/text_texture_generation/lib2/init_view.py b/modelscope/models/cv/text_texture_generation/lib2/init_view.py new file mode 100644 index 00000000..5bc44e7a --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/lib2/init_view.py @@ -0,0 +1,229 @@ +PALETTE = { + 0: [255, 255, 255], # white - background + 1: [204, 50, 50], # red - old + 2: [231, 180, 22], # yellow - update + 3: [45, 201, 55] # green - new +} + +QUAD_WEIGHTS = { + 0: 0, # background + 1: 0.1, # old + 2: 0.5, # update + 3: 1 # new +} + +VIEWPOINTS = { + 2: { + 'azim': [0, 180], + 'elev': [0, 0], + 'sector': ['front', 'back'] + }, + 4: { + 'azim': [ + 45, + 315, + 135, + 225, + ], + 'elev': [ + 0, + 0, + 0, + 0, + ], + 'sector': [ + 'front right', + 'front left', + 'back right', + 'back left', + ] + }, + 6: { + 'azim': [0, 90, 270, 0, 180, 0], + 'elev': [0, 0, 0, 90, 0, -90], + 'sector': [ + 'front', + 'right', + 'left', + 'top', + 'back', + 'bottom', + ] + }, + 10: { + 'azim': [270, 315, 225, 0, 180, 45, 135, 90, 270, 270], + 'elev': [15, 15, 15, 15, 15, 15, 15, 15, 90, -90], + 'sector': [ + 'front', + 'front right', + 'front left', + 'right', + 'left', + 'back right', + 'back left', + 'back', + 'top', + 'bottom', + ] + }, + 12: { + 'azim': [ + 0, + 45, + 315, + 135, + 225, + 180, + 45, + 315, + 90, + 270, + 90, + 270, + ], + 'elev': [ + 0, + 0, + 0, + 0, + 0, + 0, + 30, + 30, + 15, + 15, + 90, + -90, + ], + 'sector': [ + 'front', + 'front right', + 'front left', + 'back right', + 'back left', + 'back', + 'front right', + 'front left', + 'right', + 'left', + 'top', + 'bottom', + ] + }, + 36: { + 'azim': [ + 45, + 315, + 135, + 225, + 0, + 45, + 315, + 90, + 270, + 135, + 225, + 180, + 0, + 45, + 315, + 90, + 270, + 135, + 225, + 180, + 22.5, + 337.5, + 67.5, + 292.5, + 112.5, + 247.5, + 157.5, + 202.5, + 22.5, + 337.5, + 67.5, + 292.5, + 112.5, + 247.5, + 157.5, + 202.5, + ], + 'elev': [ + 0, + 0, + 0, + 0, + 30, + 30, + 30, + 30, + 30, + 30, + 30, + 30, + 60, + 60, + 60, + 60, + 60, + 60, + 60, + 60, + 15, + 15, + 15, + 15, + 15, + 15, + 15, + 15, + 45, + 45, + 45, + 45, + 45, + 45, + 45, + 45, + ], + 'sector': [ + 'front right', + 'front left', + 'back right', + 'back left', + 'front', + 'front right', + 'front left', + 'right', + 'left', + 'back right', + 'back left', + 'back', + 'top front', + 'top right', + 'top left', + 'top right', + 'top left', + 'top right', + 'top left', + 'top back', + 'front right', + 'front left', + 'front right', + 'front left', + 'back right', + 'back left', + 'back right', + 'back left', + 'front right', + 'front left', + 'front right', + 'front left', + 'back right', + 'back left', + 'back right', + 'back left', + ] + } +} diff --git a/modelscope/models/cv/text_texture_generation/lib2/projection.py b/modelscope/models/cv/text_texture_generation/lib2/projection.py new file mode 100644 index 00000000..4209eae8 --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/lib2/projection.py @@ -0,0 +1,655 @@ +import os +import random +# customized +import sys +from typing import NamedTuple, Sequence + +import cv2 +import numpy as np +import torch +from PIL import Image +from pytorch3d.io import save_obj +from pytorch3d.ops import interpolate_face_attributes +from pytorch3d.renderer import (AmbientLights, MeshRasterizer, + MeshRendererWithFragments, + RasterizationSettings, SoftPhongShader, + TexturesUV) +from pytorch3d.renderer.mesh.shader import ShaderBase +from torchvision import transforms +from tqdm import tqdm + +from modelscope.models.cv.text_texture_generation.lib2.camera import \ + init_camera +from modelscope.models.cv.text_texture_generation.lib2.init_view import * +from modelscope.models.cv.text_texture_generation.lib2.viusel import ( + visualize_outputs, visualize_quad_mask) + +sys.path.append('.') + + +class BlendParams(NamedTuple): + sigma: float = 1e-4 + gamma: float = 1e-4 + background_color: Sequence = (1, 1, 1) + + +class FlatTexelShader(ShaderBase): + + def __init__(self, + device='cpu', + cameras=None, + lights=None, + materials=None, + blend_params=None): + super().__init__(device, cameras, lights, materials, blend_params) + + def forward(self, fragments, meshes, **_kwargs): + texels = meshes.sample_textures(fragments) + texels[(fragments.pix_to_face == -1), :] = 0 + return texels.squeeze(-2) + + +def init_soft_phong_shader(camera, blend_params, device): + lights = AmbientLights(device=device) + shader = SoftPhongShader( + cameras=camera, + lights=lights, + device=device, + blend_params=blend_params) + + return shader + + +def init_flat_texel_shader(camera, device): + shader = FlatTexelShader(cameras=camera, device=device) + return shader + + +def init_renderer(camera, shader, image_size, faces_per_pixel): + raster_settings = RasterizationSettings( + image_size=image_size, faces_per_pixel=faces_per_pixel) + renderer = MeshRendererWithFragments( + rasterizer=MeshRasterizer( + cameras=camera, raster_settings=raster_settings), + shader=shader) + + return renderer + + +@torch.no_grad() +def render(mesh, renderer, pad_value=10): + + def phong_normal_shading(meshes, fragments) -> torch.Tensor: + faces = meshes.faces_packed() # (F, 3) + vertex_normals = meshes.verts_normals_packed() # (V, 3) + faces_normals = vertex_normals[faces] + pixel_normals = interpolate_face_attributes(fragments.pix_to_face, + fragments.bary_coords, + faces_normals) + + return pixel_normals + + def similarity_shading(meshes, fragments): + faces = meshes.faces_packed() # (F, 3) + vertex_normals = meshes.verts_normals_packed() # (V, 3) + faces_normals = vertex_normals[faces] + vertices = meshes.verts_packed() # (V, 3) + face_positions = vertices[faces] + view_directions = torch.nn.functional.normalize( + (renderer.shader.cameras.get_camera_center().reshape(1, 1, 3) + - face_positions), + p=2, + dim=2) + cosine_similarity = torch.nn.CosineSimilarity(dim=2)(faces_normals, + view_directions) + pixel_similarity = interpolate_face_attributes( + fragments.pix_to_face, fragments.bary_coords, + cosine_similarity.unsqueeze(-1)) + + return pixel_similarity + + def get_relative_depth_map(fragments, pad_value=pad_value): + absolute_depth = fragments.zbuf[..., 0] # B, H, W + no_depth = -1 + + depth_min, depth_max = absolute_depth[absolute_depth != no_depth].min( + ), absolute_depth[absolute_depth != no_depth].max() + target_min, target_max = 50, 255 + + depth_value = absolute_depth[absolute_depth != no_depth] + depth_value = depth_max - depth_value # reverse values + + depth_value /= (depth_max - depth_min) + depth_value = depth_value * (target_max - target_min) + target_min + + relative_depth = absolute_depth.clone() + relative_depth[absolute_depth != no_depth] = depth_value + relative_depth[absolute_depth == no_depth] = pad_value + + return relative_depth + + images, fragments = renderer(mesh) + normal_maps = phong_normal_shading(mesh, fragments).squeeze(-2) + similarity_maps = similarity_shading(mesh, fragments).squeeze(-2) # -1 - 1 + depth_maps = get_relative_depth_map(fragments) + + # normalize similarity mask to 0 - 1 + similarity_maps = torch.abs(similarity_maps) # 0 - 1 + + # HACK erode, eliminate isolated dots + non_zero_similarity = (similarity_maps > 0).float() + non_zero_similarity = (non_zero_similarity * 255.).cpu().numpy().astype( + np.uint8)[0] + non_zero_similarity = cv2.erode( + non_zero_similarity, kernel=np.ones((3, 3), np.uint8), iterations=2) + non_zero_similarity = torch.from_numpy(non_zero_similarity).to( + similarity_maps.device).unsqueeze(0) / 255. + similarity_maps = non_zero_similarity.unsqueeze(-1) * similarity_maps + return images, normal_maps, similarity_maps, depth_maps, fragments + + +@torch.no_grad() +def check_visible_faces(mesh, fragments): + pix_to_face = fragments.pix_to_face + visible_map = pix_to_face.unique() # (num_visible_faces) + return visible_map + + +def get_all_4_locations(values_y, values_x): + y_0 = torch.floor(values_y) + y_1 = torch.ceil(values_y) + x_0 = torch.floor(values_x) + x_1 = torch.ceil(values_x) + + return torch.cat([y_0, y_0, y_1, y_1], + 0).long(), torch.cat([x_0, x_1, x_0, x_1], 0).long() + + +def compose_quad_mask(new_mask_image, update_mask_image, old_mask_image, + device): + """ + compose quad mask: + -> 0: background + -> 1: old + -> 2: update + -> 3: new + """ + + new_mask_tensor = transforms.ToTensor()(new_mask_image).to(device) + update_mask_tensor = transforms.ToTensor()(update_mask_image).to(device) + old_mask_tensor = transforms.ToTensor()(old_mask_image).to(device) + + all_mask_tensor = new_mask_tensor + update_mask_tensor + old_mask_tensor + + quad_mask_tensor = torch.zeros_like(all_mask_tensor) + quad_mask_tensor[old_mask_tensor == 1] = 1 + quad_mask_tensor[update_mask_tensor == 1] = 2 + quad_mask_tensor[new_mask_tensor == 1] = 3 + + return old_mask_tensor, update_mask_tensor, new_mask_tensor, all_mask_tensor, quad_mask_tensor + + +def compute_view_heat(similarity_tensor, quad_mask_tensor): + num_total_pixels = quad_mask_tensor.reshape(-1).shape[0] + heat = 0 + for idx in QUAD_WEIGHTS: + heat += (quad_mask_tensor + == idx).sum() * QUAD_WEIGHTS[idx] / num_total_pixels + + return heat + + +def select_viewpoint(selected_view_ids, + view_punishments, + mode, + dist_list, + elev_list, + azim_list, + sector_list, + view_idx, + similarity_texture_cache, + exist_texture, + mesh, + faces, + verts_uvs, + image_size, + faces_per_pixel, + init_image_dir, + mask_image_dir, + normal_map_dir, + depth_map_dir, + similarity_map_dir, + device, + use_principle=False): + if mode == 'sequential': + + num_views = len(dist_list) + + dist = dist_list[view_idx % num_views] + elev = elev_list[view_idx % num_views] + azim = azim_list[view_idx % num_views] + sector = sector_list[view_idx % num_views] + + selected_view_ids.append(view_idx % num_views) + + elif mode == 'heuristic': + + if use_principle and view_idx < 6: + + selected_view_idx = view_idx + + else: + + selected_view_idx = None + max_heat = 0 + + print('=> selecting next view...') + view_heat_list = [] + for sample_idx in tqdm(range(len(dist_list))): + + view_heat, *_ = render_one_view_and_build_masks( + dist_list[sample_idx], elev_list[sample_idx], + azim_list[sample_idx], sample_idx, sample_idx, + view_punishments, similarity_texture_cache, exist_texture, + mesh, faces, verts_uvs, image_size, faces_per_pixel, + init_image_dir, mask_image_dir, normal_map_dir, + depth_map_dir, similarity_map_dir, device) + + if view_heat > max_heat: + selected_view_idx = sample_idx + max_heat = view_heat + + view_heat_list.append(view_heat.item()) + + print(view_heat_list) + print('select view {} with heat {}'.format(selected_view_idx, + max_heat)) + + dist = dist_list[selected_view_idx] + elev = elev_list[selected_view_idx] + azim = azim_list[selected_view_idx] + sector = sector_list[selected_view_idx] + + selected_view_ids.append(selected_view_idx) + + view_punishments[selected_view_idx] *= 0.01 + + elif mode == 'random': + + selected_view_idx = random.choice(range(len(dist_list))) + + dist = dist_list[selected_view_idx] + elev = elev_list[selected_view_idx] + azim = azim_list[selected_view_idx] + sector = sector_list[selected_view_idx] + + selected_view_ids.append(selected_view_idx) + + else: + raise NotImplementedError() + + return dist, elev, azim, sector, selected_view_ids, view_punishments + + +@torch.no_grad() +def build_backproject_mask(mesh, faces, verts_uvs, cameras, reference_image, + faces_per_pixel, image_size, uv_size, device): + # construct pixel UVs + renderer_scaled = init_renderer( + cameras, + shader=init_soft_phong_shader( + camera=cameras, blend_params=BlendParams(), device=device), + image_size=image_size, + faces_per_pixel=faces_per_pixel) + fragments_scaled = renderer_scaled.rasterizer(mesh) + + # get UV coordinates for each pixel + faces_verts_uvs = verts_uvs[faces.textures_idx] + + pixel_uvs = interpolate_face_attributes(fragments_scaled.pix_to_face, + fragments_scaled.bary_coords, + faces_verts_uvs) # NxHsxWsxKx2 + pixel_uvs = pixel_uvs.permute(0, 3, 1, 2, 4).reshape(-1, 2) + + texture_locations_y, texture_locations_x = get_all_4_locations( + (1 - pixel_uvs[:, 1]).reshape(-1) * (uv_size - 1), + pixel_uvs[:, 0].reshape(-1) * (uv_size - 1)) + + K = faces_per_pixel + + texture_values = torch.from_numpy( + np.array(reference_image.resize( + (image_size, image_size)))).float() / 255. + texture_values = texture_values.to(device).unsqueeze(0).expand( + [4, -1, -1, -1]).unsqueeze(0).expand([K, -1, -1, -1, -1]) + + # texture + texture_tensor = torch.zeros(uv_size, uv_size, 3).to(device) + texture_tensor[texture_locations_y, + texture_locations_x, :] = texture_values.reshape(-1, 3) + + return texture_tensor[:, :, 0] + + +@torch.no_grad() +def build_diffusion_mask(mesh_stuff, + renderer, + exist_texture, + similarity_texture_cache, + target_value, + device, + image_size, + smooth_mask=False, + view_threshold=0.01): + mesh, faces, verts_uvs = mesh_stuff + mask_mesh = mesh.clone() # NOTE in-place operation - DANGER!!! + + # visible mask => the whole region + exist_texture_expand = exist_texture.unsqueeze(0).unsqueeze(-1).expand( + -1, -1, -1, 3).to(device) + mask_mesh.textures = TexturesUV( + maps=torch.ones_like(exist_texture_expand), + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=verts_uvs[None, ...], + sampling_mode='nearest') + # visible_mask_tensor, *_ = render(mask_mesh, renderer) + visible_mask_tensor, _, similarity_map_tensor, *_ = render( + mask_mesh, renderer) + # faces that are too rotated away from the viewpoint will be treated as invisible + valid_mask_tensor = (similarity_map_tensor >= view_threshold).float() + visible_mask_tensor *= valid_mask_tensor + + # nonexist mask <=> new mask + exist_texture_expand = exist_texture.unsqueeze(0).unsqueeze(-1).expand( + -1, -1, -1, 3).to(device) + mask_mesh.textures = TexturesUV( + maps=1 - exist_texture_expand, + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=verts_uvs[None, ...], + sampling_mode='nearest') + new_mask_tensor, *_ = render(mask_mesh, renderer) + new_mask_tensor *= valid_mask_tensor + + # exist mask => visible mask - new mask + exist_mask_tensor = visible_mask_tensor - new_mask_tensor + exist_mask_tensor[ + exist_mask_tensor < 0] = 0 # NOTE dilate can lead to overflow + + # all update mask + mask_mesh.textures = TexturesUV( + maps=( + similarity_texture_cache.argmax(0) == target_value + # # only consider the views that have already appeared before + # similarity_texture_cache[0:target_value+1].argmax(0) == target_value + ).float().unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3).to(device), + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=verts_uvs[None, ...], + sampling_mode='nearest') + all_update_mask_tensor, *_ = render(mask_mesh, renderer) + + # current update mask => intersection between all update mask and exist mask + update_mask_tensor = exist_mask_tensor * all_update_mask_tensor + + # keep mask => exist mask - update mask + old_mask_tensor = exist_mask_tensor - update_mask_tensor + + # convert + new_mask = new_mask_tensor[0].cpu().float().permute(2, 0, 1) + new_mask = transforms.ToPILImage()(new_mask).convert('L') + + update_mask = update_mask_tensor[0].cpu().float().permute(2, 0, 1) + update_mask = transforms.ToPILImage()(update_mask).convert('L') + + old_mask = old_mask_tensor[0].cpu().float().permute(2, 0, 1) + old_mask = transforms.ToPILImage()(old_mask).convert('L') + + exist_mask = exist_mask_tensor[0].cpu().float().permute(2, 0, 1) + exist_mask = transforms.ToPILImage()(exist_mask).convert('L') + + return new_mask, update_mask, old_mask, exist_mask + + +@torch.no_grad() +def render_one_view(mesh, dist, elev, azim, image_size, faces_per_pixel, + device): + # render the view + # print(image_size) + cameras = init_camera(dist, elev, azim, image_size, device) + renderer = init_renderer( + cameras, + shader=init_soft_phong_shader( + camera=cameras, blend_params=BlendParams(), device=device), + image_size=image_size, + faces_per_pixel=faces_per_pixel) + + init_images_tensor, normal_maps_tensor, similarity_tensor, depth_maps_tensor, fragments = render( + mesh, renderer) + # print(init_images_tensor.shape, torch.max(init_images_tensor), torch.min(init_images_tensor)) + cv2.imwrite('img.png', + (np.array(init_images_tensor.squeeze(0)[:, :, :3].cpu()) + * 255).astype(np.uint8)) + return (cameras, renderer, init_images_tensor, normal_maps_tensor, + similarity_tensor, depth_maps_tensor, fragments) + + +@torch.no_grad() +def build_similarity_texture_cache_for_all_views(mesh, faces, verts_uvs, + dist_list, elev_list, + azim_list, image_size, + image_size_scaled, uv_size, + faces_per_pixel, device): + num_candidate_views = len(dist_list) + similarity_texture_cache = torch.zeros(num_candidate_views, uv_size, + uv_size).to(device) + + print('=> building similarity texture cache for all views...') + for i in tqdm(range(num_candidate_views)): + cameras, _, _, _, similarity_tensor, _, _ = render_one_view( + mesh, dist_list[i], elev_list[i], azim_list[i], image_size, + faces_per_pixel, device) + + similarity_texture_cache[i] = build_backproject_mask( + mesh, faces, verts_uvs, cameras, + transforms.ToPILImage()(similarity_tensor[0, :, :, + 0]).convert('RGB'), + faces_per_pixel, image_size_scaled, uv_size, device) + + return similarity_texture_cache + + +@torch.no_grad() +def render_one_view_and_build_masks(dist, + elev, + azim, + selected_view_idx, + view_idx, + view_punishments, + similarity_texture_cache, + exist_texture, + mesh, + faces, + verts_uvs, + image_size, + faces_per_pixel, + init_image_dir, + mask_image_dir, + normal_map_dir, + depth_map_dir, + similarity_map_dir, + device, + save_intermediate=False, + smooth_mask=False, + view_threshold=0.01): + # render the view + (cameras, renderer, init_images_tensor, normal_maps_tensor, + similarity_tensor, depth_maps_tensor, + fragments) = render_one_view(mesh, dist, elev, azim, image_size, + faces_per_pixel, device) + + init_image = init_images_tensor[0].cpu() + init_image = init_image.permute(2, 0, 1) + init_image = transforms.ToPILImage()(init_image).convert('RGB') + + normal_map = normal_maps_tensor[0].cpu() + normal_map = normal_map.permute(2, 0, 1) + normal_map = transforms.ToPILImage()(normal_map).convert('RGB') + + depth_map = depth_maps_tensor[0].cpu().numpy() + depth_map = Image.fromarray(depth_map).convert('L') + + similarity_map = similarity_tensor[0, :, :, 0].cpu() + similarity_map = transforms.ToPILImage()(similarity_map).convert('L') + + flat_renderer = init_renderer( + cameras, + shader=init_flat_texel_shader(camera=cameras, device=device), + image_size=image_size, + faces_per_pixel=faces_per_pixel) + new_mask_image, update_mask_image, old_mask_image, exist_mask_image = build_diffusion_mask( + (mesh, faces, verts_uvs), + flat_renderer, + exist_texture, + similarity_texture_cache, + selected_view_idx, + device, + image_size, + smooth_mask=smooth_mask, + view_threshold=view_threshold) + # NOTE the view idx is the absolute idx in the sample space (i.e. `selected_view_idx`) + # it should match with `similarity_texture_cache` + + (old_mask_tensor, update_mask_tensor, new_mask_tensor, all_mask_tensor, + quad_mask_tensor) = compose_quad_mask(new_mask_image, update_mask_image, + old_mask_image, device) + + view_heat = compute_view_heat(similarity_tensor, quad_mask_tensor) + view_heat *= view_punishments[selected_view_idx] + + # save intermediate results + if save_intermediate: + init_image.save( + os.path.join(init_image_dir, '{}.png'.format(view_idx))) + normal_map.save( + os.path.join(normal_map_dir, '{}.png'.format(view_idx))) + depth_map.save(os.path.join(depth_map_dir, '{}.png'.format(view_idx))) + similarity_map.save( + os.path.join(similarity_map_dir, '{}.png'.format(view_idx))) + + new_mask_image.save( + os.path.join(mask_image_dir, '{}_new.png'.format(view_idx))) + update_mask_image.save( + os.path.join(mask_image_dir, '{}_update.png'.format(view_idx))) + old_mask_image.save( + os.path.join(mask_image_dir, '{}_old.png'.format(view_idx))) + exist_mask_image.save( + os.path.join(mask_image_dir, '{}_exist.png'.format(view_idx))) + + visualize_quad_mask(mask_image_dir, quad_mask_tensor, view_idx, + view_heat, device) + + return (view_heat, renderer, cameras, fragments, init_image, normal_map, + depth_map, init_images_tensor, normal_maps_tensor, + depth_maps_tensor, similarity_tensor, old_mask_image, + update_mask_image, new_mask_image, old_mask_tensor, + update_mask_tensor, new_mask_tensor, all_mask_tensor, + quad_mask_tensor) + + +def save_full_obj(output_dir, obj_name, verts, faces, verts_uvs, faces_uvs, + projected_texture, device): + print('=> saving OBJ file...') + texture_map = transforms.ToTensor()(projected_texture).to(device) + texture_map = texture_map.permute(1, 2, 0) + obj_path = os.path.join(output_dir, obj_name) + + save_obj( + obj_path, + verts=verts, + faces=faces, + decimal_places=5, + verts_uvs=verts_uvs, + faces_uvs=faces_uvs, + texture_map=texture_map) + + +@torch.no_grad() +def backproject_from_image(mesh, faces, verts_uvs, cameras, reference_image, + new_mask_image, update_mask_image, init_texture, + exist_texture, image_size, uv_size, faces_per_pixel, + device): + # construct pixel UVs + renderer_scaled = init_renderer( + cameras, + shader=init_soft_phong_shader( + camera=cameras, blend_params=BlendParams(), device=device), + image_size=image_size, + faces_per_pixel=faces_per_pixel) + fragments_scaled = renderer_scaled.rasterizer(mesh) + + # get UV coordinates for each pixel + faces_verts_uvs = verts_uvs[faces.textures_idx] + + pixel_uvs = interpolate_face_attributes(fragments_scaled.pix_to_face, + fragments_scaled.bary_coords, + faces_verts_uvs) # NxHsxWsxKx2 + pixel_uvs = pixel_uvs.permute(0, 3, 1, 2, + 4).reshape(pixel_uvs.shape[-2], + pixel_uvs.shape[1], + pixel_uvs.shape[2], 2) + + # the update mask has to be on top of the diffusion mask + new_mask_image_tensor = transforms.ToTensor()(new_mask_image).to( + device).unsqueeze(-1) + update_mask_image_tensor = transforms.ToTensor()(update_mask_image).to( + device).unsqueeze(-1) + + project_mask_image_tensor = torch.logical_or( + update_mask_image_tensor, new_mask_image_tensor).float() + project_mask_image = project_mask_image_tensor * 255. + project_mask_image = Image.fromarray( + project_mask_image[0, :, :, 0].cpu().numpy().astype(np.uint8)) + + project_mask_image_scaled = project_mask_image.resize( + (image_size, image_size), ) + # Image.Resampling.NEAREST + # ) + project_mask_image_tensor_scaled = transforms.ToTensor()( + project_mask_image_scaled).to(device) + + pixel_uvs_masked = pixel_uvs[project_mask_image_tensor_scaled == 1] + + texture_locations_y, texture_locations_x = get_all_4_locations( + (1 - pixel_uvs_masked[:, 1]).reshape(-1) * (uv_size - 1), + pixel_uvs_masked[:, 0].reshape(-1) * (uv_size - 1)) + + K = pixel_uvs.shape[0] + project_mask_image_tensor_scaled = project_mask_image_tensor_scaled[:, + None, :, :, + None].repeat( + 1, + 4, + 1, + 1, + 3) + + texture_values = torch.from_numpy( + np.array(reference_image.resize((image_size, image_size)))) + texture_values = texture_values.to(device).unsqueeze(0).expand( + [4, -1, -1, -1]).unsqueeze(0).expand([K, -1, -1, -1, -1]) + + texture_values_masked = texture_values.reshape( + -1, 3)[project_mask_image_tensor_scaled.reshape(-1, 3) == 1].reshape( + -1, 3) + + # texture + texture_tensor = torch.from_numpy(np.array(init_texture)).to(device) + texture_tensor[texture_locations_y, + texture_locations_x, :] = texture_values_masked + + init_texture = Image.fromarray(texture_tensor.cpu().numpy().astype( + np.uint8)) + + # update texture cache + exist_texture[texture_locations_y, texture_locations_x] = 1 + + return init_texture, project_mask_image, exist_texture diff --git a/modelscope/models/cv/text_texture_generation/lib2/viusel.py b/modelscope/models/cv/text_texture_generation/lib2/viusel.py new file mode 100644 index 00000000..faf466bb --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/lib2/viusel.py @@ -0,0 +1,268 @@ +import os +import sys + +import imageio.v2 as imageio +# visualization +import matplotlib +import matplotlib.cm as cm +import matplotlib.pyplot as plt +import numpy as np +import torch +from PIL import Image + +from modelscope.models.cv.text_texture_generation.lib2.camera import \ + polar_to_xyz +from modelscope.models.cv.text_texture_generation.lib2.init_view import * + +matplotlib.use('Agg') + +sys.path.append('.') + + +def visualize_quad_mask(mask_image_dir, quad_mask_tensor, view_idx, view_score, + device): + quad_mask_tensor = quad_mask_tensor.unsqueeze(-1).repeat(1, 1, 1, 3) + quad_mask_image_tensor = torch.zeros_like(quad_mask_tensor) + + for idx in PALETTE: + selected = quad_mask_tensor[quad_mask_tensor == idx].reshape(-1, 3) + selected = torch.FloatTensor( + PALETTE[idx]).to(device).unsqueeze(0).repeat(selected.shape[0], 1) + + quad_mask_image_tensor[quad_mask_tensor == idx] = selected.reshape(-1) + + quad_mask_image_np = quad_mask_image_tensor[0].cpu().numpy().astype( + np.uint8) + quad_mask_image = Image.fromarray(quad_mask_image_np).convert('RGB') + quad_mask_image.save( + os.path.join(mask_image_dir, + '{}_quad_{:.5f}.png'.format(view_idx, view_score))) + + +def visualize_outputs(output_dir, init_image_dir, mask_image_dir, + inpainted_image_dir, num_views): + # subplot settings + num_col = 3 + num_row = 1 + sus = 4 + + summary_image_dir = os.path.join(output_dir, 'summary') + os.makedirs(summary_image_dir, exist_ok=True) + + # graph settings + print('=> visualizing results...') + for view_idx in range(num_views): + plt.switch_backend('agg') + fig = plt.figure(dpi=100) + fig.set_size_inches(sus * num_col, sus * (num_row + 1)) + fig.set_facecolor('white') + + # rendering + plt.subplot2grid((num_row, num_col), (0, 0)) + plt.imshow( + Image.open( + os.path.join(init_image_dir, '{}.png'.format(view_idx)))) + plt.text( + 0, + 0, + 'Rendering', + fontsize=16, + color='black', + backgroundcolor='white') + plt.axis('off') + + # mask + plt.subplot2grid((num_row, num_col), (0, 1)) + plt.imshow( + Image.open( + os.path.join(mask_image_dir, + '{}_project.png'.format(view_idx)))) + plt.text( + 0, + 0, + 'Project Mask', + fontsize=16, + color='black', + backgroundcolor='white') + plt.set_cmap(cm.Greys_r) + plt.axis('off') + + # inpainted + plt.subplot2grid((num_row, num_col), (0, 2)) + plt.imshow( + Image.open( + os.path.join(inpainted_image_dir, '{}.png'.format(view_idx)))) + plt.text( + 0, + 0, + 'Inpainted', + fontsize=16, + color='black', + backgroundcolor='white') + plt.axis('off') + + plt.savefig( + os.path.join(summary_image_dir, '{}.png'.format(view_idx)), + bbox_inches='tight') + fig.clf() + + # generate GIF + images = [ + imageio.imread( + os.path.join(summary_image_dir, '{}.png'.format(view_idx))) + for view_idx in range(num_views) + ] + imageio.mimsave( + os.path.join(summary_image_dir, 'output.gif'), images, duration=1) + + print('=> done!') + + +def visualize_principle_viewpoints(output_dir, dist_list, elev_list, + azim_list): + theta_list = [e for e in azim_list] + phi_list = [90 - e for e in elev_list] + DIST = dist_list[0] + + xyz_list = [ + polar_to_xyz(theta, phi, DIST) + for theta, phi in zip(theta_list, phi_list) + ] + + xyz_np = np.array(xyz_list) + color_np = np.array([[0, 0, 0]]).repeat(xyz_np.shape[0], 0) + + ax = plt.axes(projection='3d') + SCALE = 0.8 + ax.set_xlim((-DIST, DIST)) + ax.set_ylim((-DIST, DIST)) + ax.set_zlim((-SCALE * DIST, SCALE * DIST)) + + ax.scatter( + xyz_np[:, 0], + xyz_np[:, 2], + xyz_np[:, 1], + s=100, + c=color_np, + depthshade=True, + label='Principle views') + ax.scatter([0], [0], [0], + c=[[1, 0, 0]], + s=100, + depthshade=True, + label='Object center') + + # draw hemisphere + # theta inclination angle + # phi azimuthal angle + n_theta = 50 # number of values for theta + n_phi = 200 # number of values for phi + r = DIST # radius of sphere + + # theta, phi = np.mgrid[0.0:0.5*np.pi:n_theta*1j, 0.0:2.0*np.pi:n_phi*1j] + theta, phi = np.mgrid[0.0:1 * np.pi:n_theta * 1j, + 0.0:2.0 * np.pi:n_phi * 1j] + + x = r * np.sin(theta) * np.cos(phi) + y = r * np.sin(theta) * np.sin(phi) + z = r * np.cos(theta) + + ax.plot_surface(x, y, z, rstride=1, cstride=1, alpha=0.25, linewidth=1) + + # Make the grid + ax.quiver( + xyz_np[:, 0], + xyz_np[:, 2], + xyz_np[:, 1], + -xyz_np[:, 0], + -xyz_np[:, 2], + -xyz_np[:, 1], + normalize=True, + length=0.3) + + ax.set_xlabel('X Label') + ax.set_ylabel('Z Label') + ax.set_zlabel('Y Label') + + ax.view_init(30, 35) + ax.legend() + + plt.show() + + plt.savefig(os.path.join(output_dir, 'principle_viewpoints.png')) + + +def visualize_refinement_viewpoints(output_dir, selected_view_ids, dist_list, + elev_list, azim_list): + theta_list = [azim_list[i] for i in selected_view_ids] + phi_list = [90 - elev_list[i] for i in selected_view_ids] + DIST = dist_list[0] + + xyz_list = [ + polar_to_xyz(theta, phi, DIST) + for theta, phi in zip(theta_list, phi_list) + ] + + xyz_np = np.array(xyz_list) + color_np = np.array([[0, 0, 0]]).repeat(xyz_np.shape[0], 0) + + fig = plt.figure() + ax = plt.axes(projection='3d') + SCALE = 0.8 + ax.set_xlim((-DIST, DIST)) + ax.set_ylim((-DIST, DIST)) + ax.set_zlim((-SCALE * DIST, SCALE * DIST)) + + ax.scatter( + xyz_np[:, 0], + xyz_np[:, 2], + xyz_np[:, 1], + c=color_np, + depthshade=True, + label='Refinement views') + ax.scatter([0], [0], [0], + c=[[1, 0, 0]], + s=100, + depthshade=True, + label='Object center') + + # draw hemisphere + # theta inclination angle + # phi azimuthal angle + n_theta = 50 # number of values for theta + n_phi = 200 # number of values for phi + r = DIST # radius of sphere + + # theta, phi = np.mgrid[0.0:0.5*np.pi:n_theta*1j, 0.0:2.0*np.pi:n_phi*1j] + theta, phi = np.mgrid[0.0:1 * np.pi:n_theta * 1j, + 0.0:2.0 * np.pi:n_phi * 1j] + + x = r * np.sin(theta) * np.cos(phi) + y = r * np.sin(theta) * np.sin(phi) + z = r * np.cos(theta) + + ax.plot_surface(x, y, z, rstride=1, cstride=1, alpha=0.25, linewidth=1) + + # Make the grid + ax.quiver( + xyz_np[:, 0], + xyz_np[:, 2], + xyz_np[:, 1], + -xyz_np[:, 0], + -xyz_np[:, 2], + -xyz_np[:, 1], + normalize=True, + length=0.3) + + ax.set_xlabel('X Label') + ax.set_ylabel('Z Label') + ax.set_zlabel('Y Label') + + ax.view_init(30, 35) + ax.legend() + + plt.show() + + plt.savefig(os.path.join(output_dir, 'refinement_viewpoints.png')) + + fig.clear() diff --git a/modelscope/models/cv/text_texture_generation/utils.py b/modelscope/models/cv/text_texture_generation/utils.py new file mode 100644 index 00000000..09c40ddd --- /dev/null +++ b/modelscope/models/cv/text_texture_generation/utils.py @@ -0,0 +1,91 @@ +# common utils +import os + +import imageio.v2 as imageio +import torch +# pytorch3d +from pytorch3d.io import load_obj, load_objs_as_meshes +from pytorch3d.renderer import (AmbientLights, MeshRasterizer, + MeshRendererWithFragments, PerspectiveCameras, + RasterizationSettings, SoftPhongShader, + look_at_view_transform) +from torchvision import transforms +from tqdm import tqdm + +IMAGE_SIZE = 768 + + +def init_mesh(model_path, device): + verts, faces, aux = load_obj(model_path, device=device) + mesh = load_objs_as_meshes([model_path], device=device) + return mesh, verts, faces, aux + + +def init_camera(num_views, dist, elev, azim, view_idx, device): + interval = 360 // num_views + azim = (azim + interval * view_idx) % 360 + R, T = look_at_view_transform(dist, elev, azim) + T[0][2] = dist + image_size = torch.tensor([IMAGE_SIZE, IMAGE_SIZE]).unsqueeze(0) + focal_length = torch.tensor(2.0) + cameras = PerspectiveCameras( + focal_length=focal_length, + R=R, + T=T, + device=device, + image_size=image_size) + return cameras, dist, elev, azim + + +def init_renderer(camera, device): + raster_settings = RasterizationSettings(image_size=IMAGE_SIZE) + lights = AmbientLights(device=device) + renderer = MeshRendererWithFragments( + rasterizer=MeshRasterizer( + cameras=camera, raster_settings=raster_settings), + shader=SoftPhongShader(cameras=camera, lights=lights, device=device)) + + return renderer + + +def generation_gif(mesh_path): + num_views = 72 + if torch.cuda.is_available(): + DEVICE = torch.device('cuda:0') + torch.cuda.set_device(DEVICE) + else: + print('no gpu avaiable') + exit() + output_dir = 'GIF-{}'.format(num_views) + os.makedirs(output_dir, exist_ok=True) + + mesh, verts, faces, aux = init_mesh(mesh_path, DEVICE) + + # rendering + print('=> rendering...') + for view_idx in tqdm(range(num_views)): + init_image_path = os.path.join(output_dir, '{}.png'.format(view_idx)) + dist = 1.8 + elev = 15 + azim = 0 + + cameras, dist, elev, azim = init_camera(num_views, dist, elev, azim, + view_idx, DEVICE) + renderer = init_renderer(cameras, DEVICE) + init_images_tensor, fragments = renderer(mesh) + + # save images + init_image = init_images_tensor[0].cpu() + init_image = init_image.permute(2, 0, 1) + init_image = transforms.ToPILImage()(init_image).convert('RGB') + init_image.save(init_image_path) + + # generate GIF + images = [ + imageio.imread(os.path.join(output_dir, '{}.png').format(v_id)) + for v_id in range(args.num_views) + ] + imageio.mimsave( + os.path.join(output_dir, 'output.gif'), images, duration=0.1) + imageio.mimsave(os.path.join(output_dir, 'output.mp4'), images, fps=25) + print('=> done!') diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 67bccbf3..0fd760eb 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -907,6 +907,14 @@ TASK_OUTPUTS = { # } Tasks.human_reconstruction: [OutputKeys.OUTPUT], + # 3D text 2 texture generation result + # { + # "output": { + # "Done" + # } + # } + Tasks.text_texture_generation: [OutputKeys.OUTPUT], + # 2D hand keypoints result for single sample # { # "keypoints": [ diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index 3a2fe03a..fdc63810 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -500,6 +500,14 @@ TASK_INPUTS = { InputType.VIDEO, Tasks.human_reconstruction: InputType.IMAGE, + Tasks.text_texture_generation: { + 'mesh_path': InputType.TEXT, + 'texture_path': InputType.TEXT, + 'prompt': InputType.TEXT, + 'uvsize': InputType.NUMBER, + 'image_size': InputType.NUMBER, + 'output_dir': InputType.NUMBER, + }, Tasks.image_reid_person: InputType.IMAGE, Tasks.video_inpainting: { diff --git a/modelscope/pipelines/cv/text_texture_generation_pipeline.py b/modelscope/pipelines/cv/text_texture_generation_pipeline.py new file mode 100644 index 00000000..896699a4 --- /dev/null +++ b/modelscope/pipelines/cv/text_texture_generation_pipeline.py @@ -0,0 +1,311 @@ +# Copyright © Alibaba, Inc. and its affiliates. +import os +import random +from typing import Any, Dict + +import numpy as np +import torch +from diffusers import (ControlNetModel, DiffusionPipeline, + EulerAncestralDiscreteScheduler, + UniPCMultistepScheduler) +from PIL import Image +from pytorch3d.renderer import TexturesUV +from torchvision import transforms + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.text_texture_generation.lib2.camera import * +from modelscope.models.cv.text_texture_generation.lib2.init_view import * +from modelscope.models.cv.text_texture_generation.lib2.projection import * +from modelscope.models.cv.text_texture_generation.lib2.viusel import * +from modelscope.models.cv.text_texture_generation.utils import * +from modelscope.outputs import OutputKeys +from modelscope.pipelines.builder import PIPELINES +from modelscope.utils.constant import Tasks + + +@PIPELINES.register_module( + Tasks.text_texture_generation, + module_name=Pipelines.text_texture_generation) +class Tex2TexturePipeline(Pipelines): + """ Stable Diffusion for text_texture_generation Pipeline. + Example: + >>> import cv2 + >>> from modelscope.outputs import OutputKeys + >>> from modelscope.pipelines import pipeline + >>> from modelscope.utils.constant import Tasks + >>> input = {'mesh_path':'data/test/mesh/mesh1.obj', 'prompt':'old backpage'} + >>> model_id = 'damo/cv_diffuser_text-texture-generation' + >>> txt2texture = pipeline(Tasks.text_texture_generation, model=model_id) + >>> output = txt2texture(input) + >>> print(output) + """ + + def __init__(self, model: str, **kwargs): + super().__init__(model=model, **kwargs) + if torch.cuda.is_available(): + self.device = torch.device('cuda') + else: + print('no gpu avaiable') + exit() + + enable_xformers_memory_efficient_attention = kwargs.get( + 'enable_xformers_memory_efficient_attention', True) + try: + if enable_xformers_memory_efficient_attention: + self.model.pipe.enable_xformers_memory_efficient_attention() + except Exception as e: + print(e) + self.model.pipe.enable_model_cpu_offload() + try: + if enable_xformers_memory_efficient_attention: + self.model.inpaintmodel.enable_xformers_memory_efficient_attention( + ) + except Exception as e: + print(e) + self.model.inpaintmodel.enable_model_cpu_offload() + + def preprocess(self, inputs) -> Dict[str, Any]: + # input: {'mesh_path':'...', 'texture_path':..., uvsize:int, updatestep:int} + mesh_path = inputs.get('mesh_path', None) + mesh, verts, faces, aux, mesh_center, scale = self.model.mesh_normalized( + mesh_path) + texture_path = inputs.get('texture_path', None) + prompt = inputs.get('prompt', 'colorful') + uvsize = inputs.get('uvsize', 1024) + image_size = inputs.get('image_size', 512) + output_dir = inputs.get('output_dir', None) + if texture_path is not None: + init_texture = Image.open(texture_path).convert('RGB').resize( + (uvsize, uvsize)) + else: + zero_map = np.ones((256, 256, 3)) * 127 + init_texture = Image.fromarray( + zero_map, model='RGB').resize((uvsize, uvsize)) + new_verts_uvs = aux.verts_uvs + mesh.textures = TexturesUV( + maps=transforms.ToTensor()(init_texture)[None, ...].permute( + 0, 2, 3, 1).to(self.device), + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=new_verts_uvs[None, ...]) + result = { + 'prompt': prompt, + 'mesh': mesh, + 'faces': faces, + 'uvsize': uvsize, + 'mesh_center': mesh_center, + 'scale': scale, + 'verts_uvs': new_verts_uvs, + 'image_size': image_size, + 'init_texture': init_texture, + 'output_dir': output_dir, + } + print('mesh load done') + return result + + def forward(self, input: Dict[str, Any]) -> Dict[str, Any]: + prompt = input['prompt'] + uvsize = input['uvsize'] + mesh = input['mesh'] + mesh_center = input['mesh_center'] + scale = input['scale'] + faces = input['faces'] + verts_uvs = input['verts_uvs'] + image_size = input['image_size'] + init_texture = input['init_texture'] + output_dir = input['output_dir'] + if output_dir is None: + output_dir = 'Gen_texture' + exist_texture = torch.from_numpy( + np.zeros([uvsize, uvsize]).astype(np.float32)).to(self.device) + + generate_dir = os.path.join(output_dir, 'generate') + os.makedirs(generate_dir, exist_ok=True) + + update_dir = os.path.join(output_dir, 'update') + os.makedirs(update_dir, exist_ok=True) + + init_image_dir = os.path.join(generate_dir, 'rendering') + os.makedirs(init_image_dir, exist_ok=True) + + normal_map_dir = os.path.join(generate_dir, 'normal') + os.makedirs(normal_map_dir, exist_ok=True) + + mask_image_dir = os.path.join(generate_dir, 'mask') + os.makedirs(mask_image_dir, exist_ok=True) + + depth_map_dir = os.path.join(generate_dir, 'depth') + os.makedirs(depth_map_dir, exist_ok=True) + + similarity_map_dir = os.path.join(generate_dir, 'similarity') + os.makedirs(similarity_map_dir, exist_ok=True) + + inpainted_image_dir = os.path.join(generate_dir, 'inpainted') + os.makedirs(inpainted_image_dir, exist_ok=True) + + mesh_dir = os.path.join(generate_dir, 'mesh') + os.makedirs(mesh_dir, exist_ok=True) + + interm_dir = os.path.join(generate_dir, 'intermediate') + os.makedirs(interm_dir, exist_ok=True) + + init_dist = 1.5 + init_elev = 10 + init_azim = 0.0 + fragment_k = 1 + (dist_list, elev_list, azim_list, sector_list, + view_punishments) = init_viewpoints( + init_dist, init_elev, init_azim, use_principle=False) + pre_similarity_texture_cache = build_similarity_texture_cache_for_all_views( + mesh, faces, verts_uvs, dist_list, elev_list, azim_list, + image_size, image_size * 8, uvsize, fragment_k, self.device) + for idx in range(len(dist_list)): + print('=> processing view {}...'.format(idx)) + dist, elev, azim, sector = dist_list[idx], elev_list[ + idx], azim_list[idx], sector_list[idx] + prompt_view = ' the {} view of {}'.format(sector, prompt) + ( + view_score, + renderer, + cameras, + fragments, + init_image, + normal_map, + depth_map, + init_images_tensor, + normal_maps_tensor, + depth_maps_tensor, + similarity_tensor, + keep_mask_image, + update_mask_image, + generate_mask_image, + keep_mask_tensor, + update_mask_tensor, + generate_mask_tensor, + all_mask_tensor, + quad_mask_tensor, + ) = render_one_view_and_build_masks( + dist, + elev, + azim, + idx, + idx, + view_punishments, + # => actual view idx and the sequence idx + pre_similarity_texture_cache, + exist_texture, + mesh, + faces, + verts_uvs, + image_size, + fragment_k, + init_image_dir, + mask_image_dir, + normal_map_dir, + depth_map_dir, + similarity_map_dir, + self.device, + save_intermediate=True, + smooth_mask=False, + view_threshold=0.1) + generate_image = self.model.pipe( + prompt_view, + init_image, + generate_mask_image, + depth_maps_tensor, + strength=1.0) + init_texture, project_mask_image, exist_texture = backproject_from_image( + mesh, faces, verts_uvs, cameras, generate_image, + generate_mask_image, generate_mask_image, init_texture, + exist_texture, image_size * 8, uvsize, 1, self.device) + mesh.textures = TexturesUV( + maps=transforms.ToTensor()(init_texture)[None, ...].permute( + 0, 2, 3, 1).to(self.device), + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=verts_uvs[None, ...]) + ( + view_score, + renderer, + cameras, + fragments, + init_image, + *_, + ) = render_one_view_and_build_masks( + dist, + elev, + azim, + idx, + idx, + view_punishments, + pre_similarity_texture_cache, + exist_texture, + mesh, + faces, + verts_uvs, + image_size, + 8.0, + init_image_dir, + mask_image_dir, + normal_map_dir, + depth_map_dir, + similarity_map_dir, + self.device, + save_intermediate=False, + smooth_mask=False, + view_threshold=0.1) + if idx > 2: + diffused_image = self.model.pipe( + prompt_view, + init_image, + update_mask_image, + depth_maps_tensor, + strength=1.0) + init_texture, project_mask_image, exist_texture = backproject_from_image( + mesh, faces, verts_uvs, cameras, diffused_image, + update_mask_image, update_mask_image, init_texture, + exist_texture, image_size * 8, uvsize, 1, self.device) + # update the mesh + mesh.textures = TexturesUV( + maps=transforms.ToTensor()(init_texture)[ + None, ...].permute(0, 2, 3, 1).to(self.device), + faces_uvs=faces.textures_idx[None, ...], + verts_uvs=verts_uvs[None, ...]) + inter_images_tensor, *_ = render(mesh, renderer) + inter_image = inter_images_tensor[0].cpu() + inter_image = inter_image.permute(2, 0, 1) + inter_image = transforms.ToPILImage()(inter_image).convert('RGB') + inter_image.save(os.path.join(interm_dir, '{}.png'.format(idx))) + exist_texture_image = exist_texture * 255. + exist_texture_image = Image.fromarray( + exist_texture_image.cpu().numpy().astype( + np.uint8)).convert('L') + exist_texture_image.save( + os.path.join(mesh_dir, '{}_texture_mask.png'.format(idx))) + + mask_image = (1 - exist_texture[None, :, :, None])[0].cpu() + mask_image = mask_image.permute(2, 0, 1) + mask_image = transforms.ToPILImage()(mask_image).convert('L') + post_texture = self.model.inpaintmodel( + prompt=prompt, + image=init_image.resize((512, 512)), + mask_image=mask_image.resize((512, 512)), + height=512, + width=512).images[0].resize((uvsize, uvsize)) + diffused_image_tensor = torch.from_numpy(np.array(post_texture)).to( + self.device) + init_images_tensor = torch.from_numpy(np.array(init_image)).to( + self.device) + mask_image_tensor = 1 - exist_texture[None, :, :, None] + init_images_tensor = diffused_image_tensor * mask_image_tensor[ + 0] + init_images_tensor * (1 - mask_image_tensor[0]) + post_texture = Image.fromarray(init_images_tensor.cpu().numpy().astype( + np.uint8)).convert('RGB') + + save_full_obj(mesh_dir, 'mesh_post.obj', + scale * mesh.verts_packed() + mesh_center, + faces.verts_idx, verts_uvs, faces.textures_idx, + post_texture, self.device) + + return {OutputKeys.OUTPUT: 'Done'} + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index d8bb99fd..330abd70 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -154,6 +154,7 @@ class CVTasks(object): # 3d human reconstruction human_reconstruction = 'human-reconstruction' + text_texture_generation = 'text-texture-generation' # image quality assessment mos image_quality_assessment_mos = 'image-quality-assessment-mos' diff --git a/tests/pipelines/test_text_texture_generation.py b/tests/pipelines/test_text_texture_generation.py new file mode 100644 index 00000000..4bb6cd6f --- /dev/null +++ b/tests/pipelines/test_text_texture_generation.py @@ -0,0 +1,59 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp +import sys +import unittest + +from modelscope.hub.snapshot_download import snapshot_download +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.pipelines.base import Pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.test_utils import test_level + +sys.path.append('.') + + +@unittest.skip('For numpy compatible trimesh numpy bool') +class TextureGenerationTest(unittest.TestCase): + + def setUp(self) -> None: + self.task = Tasks.text_texture_generation + self.model_id = 'damo/cv_diffuser_text-texture-generation' + self.test_mesh = 'data/test/mesh/texture_generation/mesh1.obj' + self.prompt = 'old backpack' + + def pipeline_inference(self, pipeline: Pipeline, input_location): + result = pipeline(input_location) + mesh = result[OutputKeys.OUTPUT] + print(f'Output to {osp.abspath("mesh_post.obj")}', mesh) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_run_by_direct_model_download(self): + model_dir = snapshot_download(self.model_id) + text_texture_generation = pipeline( + Tasks.text_texture_generation, model=model_dir) + input = { + 'mesh_path': self.test_mesh, + 'prompt': self.prompt, + 'image_size': 512, + 'uvsize': 1024 + } + print('running') + self.pipeline_inference(text_texture_generation, input) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_modelhub(self): + text_texture_generation = pipeline( + Tasks.text_texture_generation, model=self.model_id) + input = { + 'mesh_path': self.test_mesh, + 'prompt': self.prompt, + 'image_size': 512, + 'uvsize': 1024 + } + print('running') + self.pipeline_inference(text_texture_generation, input) + + +if __name__ == '__main__': + unittest.main()