o ߥiVz@sddlZddlmZmZmZmZmZmZddlZddl Z ddl Z ddl m Z ddlZddlmZddlmZmZmZmZmZmZmZmZmZmZddlmZddlmZddl m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'ddl(m)Z)m*Z*m+Z+ddl,m-Z-dd l.m/Z/m0Z0dd l1m2Z2dd l3Tdd l4Tdd l5Tdd l6m7Z7m8Z8dd l9m:Z:e:Z;dZe-j>dGddde0Z? dddZ@GdddeZAdS)N)AnyCallableDictListOptionalUnion) AutoencoderKLControlNetModelDiffusionPipelineEulerAncestralDiscreteSchedulerEulerDiscreteScheduler(StableDiffusionControlNetImg2ImgPipeline(StableDiffusionControlNetInpaintPipelineStableDiffusionInpaintPipelineStableDiffusionPipelineUNet2DConditionModel)MultiControlNetModel)StableDiffusionPipelineOutput) deprecateis_accelerate_availableis_accelerate_versionis_compiled_modulelogging randn_tensorreplace_example_docstring)load_objload_objs_as_meshessave_obj)Models)Tensor TorchModel)MODELS)*) ModelFileTasks) get_loggera 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] ``` ) module_namecs@eZdZfddZddZddZd dd Zd d d ZZS) Tex2Texturecstj|d|i|tjr td|_td|jnt dt |d}|d}|d}| dtj }t j||d |j|_tj||d |j|_tj||j|d |j|_td d S) aThe 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 model_dircudaz Use GPU: {}zno gpu avaiablez /base_model/z/control_model/z/inpaint_model/ torch_dtype)r*) controlnetr*zmodel load overN)super__init__torchr) is_availabledeviceloggerinfoformatprintexitgetfloat16r from_pretrainedtor+r inpaintmodel%StableDiffusionControlinpaintPipelinepipe)selfr(argskwargs model_pathcontrolmodel_pathinpaintmodel_pathr* __class__l/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/cv/text_texture_generation/Tex2Texture.pyr-Vs8  zTex2Texture.__init__cCs0t||jd\}}}t|g|jd}||||fS)Nr0)rr0r)r= mesh_pathvertsfacesauxmeshrErErF init_meshts zTex2Texture.init_meshc Cs|}|jd}|jdd|d}|| }|ddddf|ddddf}|}d|}|d|}||}||||fS)Nrdimg?) get_bounding_boxes verts_packedshapemeanrepeat offset_vertsmax unsqueeze scale_verts) r=rLbbox num_verts mesh_centerlensmax_lenscalenew_meshrErErFnormalize_meshys $ zTex2Texture.normalize_meshnormalized.objc Cs@td|}t|||jd|j|j|jt|jdddS)Nz!=> saving normalized mesh file...r)rIrJdecimal_places verts_uvs faces_uvs texture_map)r4r verts_idxrf textures_idxtexture_imageslistkeys)r=rIrJrKpathobj_pathrErErFsave_normalized_objs zTex2Texture.save_normalized_objc CsD||\}}}}||\}}}}|||||||||||fSN)rMrbrp) r=rH save_pathrLrIrJrKr]r`rErErFmesh_normalizedszTex2Texture.mesh_normalized)rc) __name__ __module__ __qualname__r-rMrbrprs __classcell__rErErCrFr'Rs    r'Fcsf|durtd|durtdt|tjrt|tjs&tdt|d|jdkr;|jddks6Jd|d}|jdkrH|dd}|jdkr_|jdd krZ|d}n|d }|jd kri|jd ksmJd |jd d|jd dksJd |jd|jdksJd| dks| d krtd| dks| d krtdd||dk<d ||dk<|j tj d}nt|tjrtdt|dt|t jjtjfr|g}t|trt|dt jjrfdd|D}dd|D}tj|dd}nt|tr"t|dtjr"tjdd|Ddd}|ddd d}t|j tj ddd}t|t jjtjfrF|g}t|trut|dt jjrufdd|D}tjdd|Ddd}|tj d}nt|trt|dtjrtjdd|Ddd}d||dk<d ||dk<t|}||dk}|r|||fS||fS) Nz"`image` input cannot be undefined.z'`mask_image` input cannot be undefined.z,`image` is a torch.Tensor but `mask` (type: z is notrz2Image outside a batch should be of shape (3, H, W)rNrQz%Image and Mask must have 4 dimensionsz4Image and Mask must have the same spatial dimensionsz,Image and Mask must have the same batch sizez Image should be in [-1, 1] rangezMask should be in [0, 1] range?)dtypez,`mask` is a torch.Tensor but `image` (type: c"g|] }|jftjjdqS)resampleresizePILImageLANCZOS.0iheightwidthrErF z1prepare_mask_and_masked_image..cSs(g|]}t|ddddfqS)RGBNnparrayconvertrrErErFrs()axiscSsg|] }|dddfqSrqrErrErErFrsg_@?cr~rrrrrErFrrcSs*g|]}t|dddddfqS)LNrrmrErErFrs*go@cSsg|] }|ddddfqSrqrErrErErFrs) ValueError isinstancer.r TypeErrortypendimrTrYminrXr9float32rrrndarrayrl concatenate transpose from_numpyastype)imagemaskrr return_image masked_imagerErrFprepare_mask_and_masked_images                      rc2@sxeZdZeee                      d$d eee efd eej e j j fd eej e j j fdeej e j j eje ej e e j j e ejfdeedeededededeeee efdeededeeeje ejfdeej deej deej deededeeeeej gdfdedeeeefd eee efd!ef.d"d#ZdS)%r;Nr2@rQpilTr|Fpromptr mask_image control_imagerrstrengthnum_inference_stepsguidance_scalenegative_promptnum_images_per_prompteta generatorlatents prompt_embedsnegative_prompt_embeds output_type return_dictcallbackcallback_stepscross_attention_kwargscontrolnet_conditioning_scale guess_modec=Cs||||\}}||||||| ||| |dur"t|tr"d}n|dur0t|tr0t|}n|jd}|j}| dk}t|j rE|j j n|j }t|t rZt|t rZ|gt|j }t|trc|jjn|j djj}|pm|}|durx|ddnd}|j||| || |||d}t|tr|j||||| | ||j||d }n't|t rg}|D]}|j||||| | ||j||d }||q|}nJt||||d d \} }!}"|jj||d |j|||d \}#}|#dd|| }$|dk}%|jjj}&|jjj}'|'d k}(|j|| |&|||j|| ||"|$|%d |(d })|(r|)\}}*}+n|)\}}*| | |!|| |||j|| | \} },|!| | }-|j"|d}.t#|#D]\}/}0|rUt$%|gdn|}1|j&|1|0}1|ru|ru|}2|j&|2|0}2|'dd}3n|1}2|}3|j |2|0|3|||dd\}4}5|r|rdd|4D}4t$%t$(|5|5g}5|'dkrt$j%|1| |,gdd}1|j|1|0|||4|5ddd}6|r|6'd\}7}8|7| |8|7}6|jj)|6|0|fi|-ddid}|'d kr|+dd}9| dd}:|/t|#dkr|j*|9|*t$+|0g}9d|:|9|:|}|/t|#dks#|/d|jj,dkr9|.-|dur9|/|dkr9||/|0|qFWdn 1sFwYt.|drh|j/durh|j0d|j 0dt$j12|dks|jj3||jjj4ddd}|5|||j\}};n|}d};|;durd g|jd} 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`. NrQrrr`)rr lora_scale) rrr batch_sizerr0r}do_classifier_free_guidancerFT)rrG)rrr0ry)rtimestepis_strength_max return_noisereturn_image_latents)totalrN)encoder_hidden_statescontrolnet_condconditioning_scalerrcSs g|] }tt||gqSrE)r.cat zeros_like)rdrErErFr;szBStableDiffusionControlinpaintPipeline.__call__.. rO)rrdown_block_additional_residualsmid_block_additional_residualrrfinal_offload_hookcpulatent)rcSsg|]}| qSrErE)rhas_nsfwrErErFrs)rdo_denormalize)imagesnsfw_content_detected):_default_height_width check_inputsrstrrllenrT_execution_devicerr+ _orig_modrfloatnetsr configglobal_pool_conditionsr6_encode_promptprepare_control_imager}appendr scheduler set_timesteps get_timestepsrVvaelatent_channelsunet in_channelsprepare_latentsprepare_mask_latentsprepare_extra_step_kwargs progress_bar enumerater.rscale_model_inputchunkrstep add_noisetensororderupdatehasattrrr9r) empty_cachedecodescaling_factorrun_safety_checkerimage_processor postprocessoffloadr)=r=rrrrrrrrrrrrrrrrrrrrrrrrr0rr+rtext_encoder_lora_scalecontrol_imagescontrol_image_rr init_image timestepslatent_timesteprnum_channels_latentsnum_channels_unetrlatents_outputsnoise image_latentsmasked_image_latentsextra_step_kwargsrrtlatent_model_inputcontrol_model_inputcontrolnet_prompt_embedsdown_block_res_samplesmid_block_res_sample noise_prednoise_pred_uncondnoise_pred_textinit_latents_proper init_maskhas_nsfw_conceptrrErErF__call__s}                          [         z.StableDiffusionControlinpaintPipeline.__call__)NNNNNNrrrNrQrNNNNrTNrQNr|F)rtrurvr.no_gradrEXAMPLE_DOC_STRINGrrrrrr FloatTensorrrrintr GeneratorboolrrrrrErErErFr;s     r;)F)Bostypingrrrrrrcv2numpyrr PIL.Imagerr.torchvision.transforms transforms diffusersrr r r r r rrrr.diffusers.pipelines.controlnet.multicontrolnetr$diffusers.pipelines.stable_diffusionrdiffusers.utilsrrrrrrr pytorch3d.iorrrmodelscope.metainformodelscope.models.baserr modelscope.models.builderr!8modelscope.models.cv.text_texture_generation.lib2.camera;modelscope.models.cv.text_texture_generation.lib2.init_view2modelscope.models.cv.text_texture_generation.utilsmodelscope.utils.constantr#r$modelscope.utils.loggerr%r1rregister_moduletext_texture_generationr'rr;rErErErFs@   0  $   +I l