add texture generation task(文本引导纹理生成)

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14123234
* add texture generation task

* add output dir

* add input

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'),

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,

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)

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

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',
+        ]
+    }
+}

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

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()

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!')

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": [

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: {

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

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'

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()