o ۷i)@s(ddlmZddlmZddlmZddlZddlZddl m Z ddl m Z ddl mZmZdd lmZdd lmZmZdd lmZmZdd lmZmZmZmZdd lmZddlm Z ddl!m"Z"ddl#m$Z$erwddl%m&m'Z(dZ)ndZ)e*e+Z,dZ-eGdddeZ.Gddde"Z/dS)) dataclass)partial)AnyN)Image)tqdm) CLIPTextModel CLIPTokenizer)PipelineImageInput) AutoencoderKLUNet2DConditionModel) DDIMScheduler LCMScheduler) BaseOutputis_torch_xla_availableloggingreplace_example_docstring)is_scipy_available) randn_tensor)DiffusionPipeline)MarigoldImageProcessorTFaE Examples: ```py >>> import diffusers >>> import torch >>> pipe = diffusers.MarigoldDepthPipeline.from_pretrained( ... "prs-eth/marigold-depth-v1-1", variant="fp16", torch_dtype=torch.float16 ... ).to("cuda") >>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg") >>> depth = pipe(image) >>> vis = pipe.image_processor.visualize_depth(depth.prediction) >>> vis[0].save("einstein_depth.png") >>> depth_16bit = pipe.image_processor.export_depth_to_16bit_png(depth.prediction) >>> depth_16bit[0].save("einstein_depth_16bit.png") ``` c@sDeZdZUdZejejBed<dejBejBed<dejBed<dS)MarigoldDepthOutputug Output class for Marigold monocular depth prediction pipeline. Args: prediction (`np.ndarray`, `torch.Tensor`): Predicted depth maps with values in the range [0, 1]. The shape is `numimages × 1 × height × width` for `torch.Tensor` or `numimages × height × width × 1` for `np.ndarray`. uncertainty (`None`, `np.ndarray`, `torch.Tensor`): Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is `numimages × 1 × height × width` for `torch.Tensor` or `numimages × height × width × 1` for `np.ndarray`. latent (`None`, `torch.Tensor`): Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline. The shape is `numimages * numensemble × 4 × latentheight × latentwidth`. predictionN uncertaintylatent) __name__ __module__ __qualname____doc__npndarraytorchTensor__annotations__r&r&j/home/ubuntu/vllm_env/lib/python3.10/site-packages/diffusers/pipelines/marigold/pipeline_marigold_depth.pyrRs rc!s<eZdZdZdZdZ     d>dededee Bd e d e d e dBd e dBd e dBdedBdedBffdd Zdedededede de dedee efdBdejdBdejeejBdBde de defdd Zejjd?d!d"Zeee  #   $ $ #    % & & d@dededBdededBd'e de de dedee efdBdejeejBdBdejeejBdBde de d(e d)e fd*d+ZdejdejdBdejdBdededeejejff d,d-Z d.ejdejfd/d0Z!e"   & 1 2 3 4 5dAd6ejd e d e de d7e d8e#d9ed:e#d;edeejejdBffdunet->vae)depth disparityNTunetvae scheduler text_encoder tokenizerprediction_typescale_invariantshift_invariantdefault_denoising_stepsdefault_processing_resolutionc st||jvrtd|d|jd|j|||||d|j|||| | dt|ddrAs"z5MarigoldDepthPipeline.check_inputs..z;`generator` device placement is not consistent in the list.zUnsupported generator type: )r@r,rArBr7 ValueErrorr;r<r1r2r ImportError isinstancedictlist enumerater!r"r# is_tensorndimshapersizercdimmaxlatent_channelsall Generator)rErHrIrJrKrLrMrNrOrPrQrRrSactual_vae_scale_factor num_imagesWHiimgH_iW_iN_imax_orignew_Hnew_Wwhshape_expectedr&rfr' check_inputss                     z"MarigoldDepthPipeline.check_inputscCst|ds i|_nt|jtstdt|jdtdi|j}|d||d<|d||d<|dur>t|fi|S|durKtdd|i|Std) N_progress_bar_configz=`self._progress_bar_config` should be of type `dict`, but is r5descleavetotalz/Either `total` or `iterable` has to be defined.r&)hasattrrrjrkrhrcgetr)rEiterablerrrprogress_bar_configr&r&r' progress_barHs  z"MarigoldDepthPipeline.progress_barrrWr!Fmatch_input_resolution output_latent return_dictc$s&j}j}|dur j}|durj}||||||| | |  }jdurEd}j|djjddd}|j |} |d_j |||||\}}} || | |\}~jj ||dd d }g}jtd||dd d D]\}|||}||}|jd}|d|}jj||d jjjd dd D]'} tj||gd d}!j|!| |d dd}"jj|"| || dj}trtq||q{tj|dd~~~~~~~!~"tjfddtdjdDdd|sdj |d}#|d krHj||gjd dRfddt|Dt \}#tjddrFtj|#dd}#nd}#|rgj j!||d d|#durgrgj j!|#||d d}#| dkrj "|#durrj "|#}##|s|#fSt$|#dS)aA Function invoked when calling the pipeline. Args: image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`), `list[torch.Tensor]`: An input image or images used as an input for the depth estimation task. For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the same width and height. num_inference_steps (`int`, *optional*, defaults to `None`): Number of denoising diffusion steps during inference. The default value `None` results in automatic selection. ensemble_size (`int`, defaults to `1`): Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. processing_resolution (`int`, *optional*, defaults to `None`): Effective processing resolution. When set to `0`, matches the larger input image dimension. This produces crisper predictions, but may also lead to the overall loss of global context. The default value `None` resolves to the optimal value from the model config. match_input_resolution (`bool`, *optional*, defaults to `True`): When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer side of the output will equal to `processing_resolution`. resample_method_input (`str`, *optional*, defaults to `"bilinear"`): Resampling method used to resize input images to `processing_resolution`. The accepted values are: `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. resample_method_output (`str`, *optional*, defaults to `"bilinear"`): Resampling method used to resize output predictions to match the input resolution. The accepted values are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`. batch_size (`int`, *optional*, defaults to `1`): Batch size; only matters when setting `ensemble_size` or passing a tensor of images. ensembling_kwargs (`dict`, *optional*, defaults to `None`) Extra dictionary with arguments for precise ensembling control. The following options are available: - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in every pixel location, can be either `"median"` or `"mean"`. - regularizer_strength (`float`, *optional*, defaults to `0.02`): Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1. - max_iter (`int`, *optional*, defaults to `2`): Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options` argument. - tol (`float`, *optional*, defaults to `1e-3`): Alignment solver tolerance. The solver stops when the tolerance is reached. - max_res (`int`, *optional*, defaults to `None`): Resolution at which the alignment is performed; `None` matches the `processing_resolution`. latents (`torch.Tensor`, or `list[torch.Tensor]`, *optional*, defaults to `None`): Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. generator (`torch.Generator`, or `list[torch.Generator]`, *optional*, defaults to `None`): Random number generator object to ensure reproducibility. output_type (`str`, *optional*, defaults to `"np"`): Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted values are: `"np"` (numpy array) or `"pt"` (torch tensor). output_uncertainty (`bool`, *optional*, defaults to `False`): When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that the `ensemble_size` argument is set to a value above 2. output_latent (`bool`, *optional*, defaults to `False`): When enabled, the output's `latent` field contains the latent codes corresponding to the predictions within the ensemble. These codes can be saved, modified, and used for subsequent calls with the `latents` argument. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple. Examples: Returns: [`~pipelines.marigold.MarigoldDepthOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.marigold.MarigoldDepthOutput`] is returned, otherwise a `tuple` is returned where the first element is the prediction, the second element is the uncertainty (or `None`), and the third is the latent (or `None`). N do_not_padTrZ)padding max_length truncationreturn_tensorsr)rbdtyperzMarigold predictions...)rr)rbFzDiffusion steps...rr)encoder_hidden_statesrrfcs"g|] }||qSr&)decode_predictionrdr{)rN pred_latentrEr&r' (sz2MarigoldDepthPipeline.__call__..cs0g|]}j|jjfipiqSr&)ensemble_depthr1r2r)rOrSrrEr&r'r=s)is_aar!)rrr)%_execution_devicerr3r4rrCr/model_max_length input_idstor.rD preprocessprepare_latentsrepeatrrangerpr- set_timesteps timestepsr#catr+step prev_sample XLA_AVAILABLExm mark_stepappend unpad_imagereshapezipresize_antialias pt_to_numpymaybe_free_model_hooksr)$rErHrIrJrKrrLrMrNrOrPrQrRrSrrrbrrxprompt text_inputstext_input_idsroriginal_resolution image_latentbatch_empty_text_embedding pred_latentsr{batch_image_latentbatch_pred_latenteffective_batch_sizetextt batch_latentnoiserr&)rNrOrSrrrEr'__call__[sZ                zMarigoldDepthPipeline.__call__cs~ddtjfddtdjdDdd}|jjj}|j|dd}|}|dur;t|j||j |j d}||fS)NcSs,t|dr |jSt|dr|jStd)N latent_distrPz3Could not access latents of provided encoder_output)rrmoderPAttributeError)encoder_outputr&r&r'retrieve_latentsus   z?MarigoldDepthPipeline.prepare_latents..retrieve_latentsc s(g|]}j||qSr&)r,encoderrNrHrrEr&r'r~sz9MarigoldDepthPipeline.prepare_latents..rr)rQrbr) r#rrrpr,rAscaling_factorrepeat_interleaverrbr)rErHrPrQrJrNrrr&rr'rms$z%MarigoldDepthPipeline.prepare_latentsrcCs|dks|jd|jjjkrtd|jjjd|jd|jj||jjjddd}|jdd d }t |d d }|d d }|S)Nr_rz Expecting 4D tensor of shape [B,z ,H,W]; got r5F)rrTrrkeepdimg?g@) rrrpr,rArtrhdecoderr]r#clip)rErrr&r&r'rs  z'MarigoldDepthPipeline.decode_predictionr^{Gz?rMbP?r)r[regularizer_strengthmax_itertolmax_resc  s|dks |jddkrtd|jddvr"tdd s* r*tddtjf fd d dtjd tjd tjf fd d ddtjdtd ttjtjdBfffdd d tjdtjd t ffdd dtjf fdd } p } |jd| r| |} || }||d\}} | } r r| }n rd}ntd| |j dd}|||}|r| |} || fS)a Ensembles the depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The alignment happens when the predictions have one or more degrees of freedom, that is when they are either affine-invariant (`scale_invariant=True` and `shift_invariant=True`), or just scale-invariant (only `scale_invariant=True`). For absolute predictions (`scale_invariant=False` and `shift_invariant=False`) alignment is skipped and only ensembling is performed. Args: depth (`torch.Tensor`): Input ensemble depth maps. scale_invariant (`bool`, *optional*, defaults to `True`): Whether to treat predictions as scale-invariant. shift_invariant (`bool`, *optional*, defaults to `True`): Whether to treat predictions as shift-invariant. output_uncertainty (`bool`, *optional*, defaults to `False`): Whether to output uncertainty map. reduction (`str`, *optional*, defaults to `"median"`): Reduction method used to ensemble aligned predictions. The accepted values are: `"mean"` and `"median"`. regularizer_strength (`float`, *optional*, defaults to `0.02`): Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1. max_iter (`int`, *optional*, defaults to `2`): Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options` argument. tol (`float`, *optional*, defaults to `1e-3`): Alignment solver tolerance. The solver stops when the tolerance is reached. max_res (`int`, *optional*, defaults to `1024`): Resolution at which the alignment is performed; `None` matches the `processing_resolution`. Returns: A tensor of aligned and ensembled depth maps and optionally a tensor of uncertainties of the same shape: `(1, 1, H, W)`. r_rz,Expecting 4D tensor of shape [B,1,H,W]; got r5r\Unrecognized reduction method: z1Pure shift-invariant ensembling is not supported.r)cs|djddj}|djddj}r5r5d||jdd}| |}t||f}nrFd|jdd}|}nt d| t j }|S)Nrrrư>minUnrecognized alignment.) rrvaluesrsclampr#rcpunumpyrhastyper!float64)r)init_mininit_maxinit_sinit_tparamrJr1r2r&r' init_params  z8MarigoldDepthPipeline.ensemble_depth..init_paramrrTcsr0r0t|d\}}t||ddd}t||ddd}|||}|SrFt||ddd}||}|Std)Nrrr)r!splitr# from_numpyrviewrh)r)rsroutrr&r'aligns z3MarigoldDepthPipeline.ensemble_depth..alignF depth_alignedreturn_uncertaintyNcsd}dkrtj|ddd}|rtj|ddd}||fSdkr=tj|dddj}|r9tjt||dddj}||fStdd)Nr]rTrr^rr5)r#r]stdr^rabsrh)rrrr)r[r&r'ensemblesz6MarigoldDepthPipeline.ensemble_depth..ensemblec sd}||}ttD]\}}||||}||d7}qdkrN|dd\}}|} d|} || | 7}|S)NgrrFrr) r# combinationsaranger]sqrtitemrrrs) rr)costrr{jdiffr_err_nearerr_far)rrrJrr&r'cost_fns z5MarigoldDepthPipeline.ensemble_depth..cost_fncsrddl}|tj}dur t|jddkr t|d}|}|jj t |d|dddd}|j S) NrrrV)r)BFGSF)maxiterdisp)methodroptions) scipyrr#float32rsrprresize_to_max_edgeoptimizeminimizerx)r)r depth_to_alignrres)rrrrrr&r' compute_params  z;MarigoldDepthPipeline.ensemble_depth..compute_paramrrrrr)F) rrrprhr#r$r!r"booltuplefloatrsrr)r)r1r2rSr[rrrrrrequires_aligningrr depth_max depth_min depth_ranger&) rrrrJrrrr[rr1r2rr'rsF.$$    z$MarigoldDepthPipeline.ensemble_depth)NTTNN)NNNT)NrNTrWrWrNNNr!FFT)TTFr^rrrr)&rrrr model_cpu_offload_seqr:r r r rrrstrrintr9r rkrr#r$rvrlrcompilerdisablerno_gradrEXAMPLE_DOC_STRINGrrrr staticmethodrr __classcell__r&r&rFr'r(hs2)    /             %  r()0 dataclassesr functoolsrtypingrrr!r#PILr tqdm.autor transformersrrrDr modelsr r schedulersr rutilsrrrrutils.import_utilsrutils.torch_utilsrpipeline_utilsrmarigold_image_processingrtorch_xla.core.xla_modelcore xla_modelrr get_loggerrr;r"rr(r&r&r&r's2