o 0 i&@sLddlmZddlZddlZddlmZmZddlmZm Z m Z m Z m Z ddl mZddlmZddlmZddlZddlmZddlmmZddlmZdd lmZdd lmZdd lm Z dd l!m"Z"m#Z#dd l$m%Z%ddl&m'Z'm(Z(ddl)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/m0Z0ddl1m2Z2m3Z3ddl4m5Z5m6Z6ddl7m8Z8m9Z9m:Z:m;Z;mm?Z?m@Z@mAZAe=ddd\ZBZCeeDZEgdZFGdddeZGGdddeGZHGdddeGZIGdddeGZJGdd d eJZKGd!d"d"eGZLGd#d$d$eJZMGd%d&d&eGZNGd'd(d(eNZOGd)d*d*eNZPGd+d,d,ePZQGd-d.d.ejRZSdS)/) annotationsN)ABCabstractmethod)CallableIterableIteratorMappingSequence)partial)locate)Any) get_logger)decollate_batch) MetaTensor) ThreadBuffer) AvgMergerMerger)Splitter)compute_importance_mapsliding_window_inference)VQVAE AutoencoderKL ControlNetDecoderOnlyTransformerDiffusionModelUNetSPADEAutoencoderKLSPADEDiffusionModelUNet)RFlowScheduler Scheduler)CenterSpatialCrop SpatialPad) BlendModeOrdering PatchKeysPytorchPadMode ensure_tupleoptional_import)CAMGradCAM GradCAMpptqdmname)Inferer PatchInferer SimpleInfererSlidingWindowInfererSaliencyInferer SliceInfererSlidingWindowInfererAdaptc@seZdZdZed d d Zd S)r-a A base class for model inference. Extend this class to support operations during inference, e.g. a sliding window method. Example code:: device = torch.device("cuda:0") transform = Compose([ToTensor(), LoadImage(image_only=True)]) data = transform(img_path).to(device) model = UNet(...).to(device) inferer = SlidingWindowInferer(...) model.eval() with torch.no_grad(): pred = inferer(inputs=data, network=model) ... inputs torch.Tensornetworkrargsr kwargsreturncOstd|jjd)a Run inference on `inputs` with the `network` model. Args: inputs: input of the model inference. network: model for inference. args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. Raises: NotImplementedError: When the subclass does not override this method. z Subclass z must implement this method.)NotImplementedError __class____name__selfr4r6r7r8r?X/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/inferers/inferer.py__call__RszInferer.__call__N) r4r5r6rr7r r8r r9r )r< __module__ __qualname____doc__rrAr?r?r?r@r->sr-c@sleZdZdZdeddddddfd5ddZd6ddZd7d!d"Zd8d)d*Zd+d,Z d-d.Z d/d0Z d9d3d4Z dS):r.a& Inference on patches instead of the whole image based on Splitter and Merger. This splits the input image into patches and then merge the resulted patches. Args: splitter: a `Splitter` object that split the inputs into patches. Defaults to None. If not provided or None, the inputs are considered to be already split into patches. In this case, the output `merged_shape` and the optional `cropped_shape` cannot be inferred and should be explicitly provided. merger_cls: a `Merger` subclass that can be instantiated to merges patch outputs. It can also be a string that matches the name of a class inherited from `Merger` class. Defaults to `AvgMerger`. batch_size: batch size for patches. If the input tensor is already batched [BxCxWxH], this adds additional batching [(Bp*B)xCxWpxHp] for inference on patches. Defaults to 1. preprocessing: a callable that process patches before the being fed to the network. Defaults to None. postprocessing: a callable that process the output of the network. Defaults to None. output_keys: if the network output is a dictionary, this defines the keys of the output dictionary to be used for merging. Defaults to None, where all the keys are used. match_spatial_shape: whether to crop the output to match the input shape. Defaults to True. buffer_size: number of patches to be held in the buffer with a separate thread for batch sampling. Defaults to 0. merger_kwargs: arguments to be passed to `merger_cls` for instantiation. `merged_shape` is calculated automatically based on the input shape and the output patch shape unless it is passed here. NTrsplitterSplitter | None merger_clstype[Merger] | str batch_sizeint preprocessingCallable | Nonepostprocessing output_keysSequence | Nonematch_spatial_shapebool buffer_size merger_kwargsr r9Nonec Ks0t|t|ttdfst|tstdt|d||_t|trAtd|d\} } | s3t |} | dur?t d|d| }t |t sNtd|d||_ | |_|durft|sftdt|d||_|dur{t|s{td t|d||_|d krt d |d||_||_||_||_dS) Nz'splitter' should be a `Splitter` object that returns: an iterable of pairs of (patch, location) or a MetaTensor that has `PatchKeys.LOCATION` metadata).z is given.zmonai.inferers.mergerr+zThe requested `merger_cls` ['z'] does not exist.z+'merger' should be a subclass of `Merger`, z-'preprocessing' should be a callable object, z.'postprocessing' should be a callable object, rEz(`batch_size` must be a positive number, )r-__init__ isinstancertype TypeErrorrFstrr&r ValueError issubclassrrHrTcallablerLrNrJrOrQrS) r>rFrHrJrLrNrOrQrSrTZvalid_merger_clsZ merger_foundr?r?r@rVs@     zPatchInferer.__init__patches9Iterable[tuple[torch.Tensor, Sequence[int]]] | MetaTensor,Iterator[tuple[torch.Tensor, Sequence, int]]c cs0t|tr4t|}td||jD] }t|j||}||||||||jtj|fVqdS|j dkrBt ||j dd}n|}dg|j}dg|j}d}|D].} | d||<| d||<|d7}||jkrt |||fVdg|j}dg|j}d}qT|dkrt |d|||fVdSdS)zGenerate batch of patches and locations Args: patches: a tensor or list of tensors Yields: A batch of patches (torch.Tensor or MetaTensor), a sequence of location tuples, and the batch size rg?)rStimeoutNrE) rWrlenrangerJminmetar#LOCATIONrSrtorchcat) r>r^ total_sizeirJbufferZ patch_batchZlocation_batch idx_in_batchsampler?r?r@_batch_samplers4 .        zPatchInferer._batch_sampleroutputstuplecsFttr|jdurt|_tfdd|jDStddS)Nc3s|]}|VqdSNr?).0kror?r@ sz5PatchInferer._ensure_tuple_outputs..T) wrap_array)rWdictrOlistkeysrpr%)r>ror?rtr@_ensure_tuple_outputss   z"PatchInferer._ensure_tuple_outputsr6rpatchr5r7r8cOs@|jr||}||g|Ri|}|jr||}||Srq)rLrNrz)r>r6r{r7r8ror?r?r@_run_inferences    zPatchInferer._run_inferencecCst||d}g}g}|D]Z}t||d} tddt|jdd| jddD} |j} ||| | \} } d| vrN| | d<| ddurNtdd| vrV| | d<|j di| }| || | q||fS) Nrcss|] \}}||VqdSrqr?)rripopr?r?r@rusz3PatchInferer._initialize_mergers.. merged_shapez `merged_shape` cannot be `None`. cropped_shaper?) rgchunkrpzipshaperTcopy_get_merged_shapesr[rHappend)r>r4ror^rJin_patchmergersratiosZout_patch_batch out_patchratiorTrrmergerr?r?r@_initialize_mergerss$,    z PatchInferer._initialize_mergersc CsXt|||D]#\}}}t|t||D]\} } ddt| |D} || | qqdS)NcSsg|] \}}t||qSr?round)rrlrr?r?r@ sz+PatchInferer._aggregate..)rrgr aggregate) r>ro locationsrJrrZoutput_patchesrrZin_locrZout_locr?r?r@ _aggregates zPatchInferer._aggregatec Cs|jdurdS|j|}|j|}tddt||D}tddt||D}|jdd|}|jdd|} |jsB| }|| fS)z:Define the shape of merged tensors (non-padded and padded)N)NNcs |] \}}t||VqdSrqrrrsrr?r?r@ru)z2PatchInferer._get_merged_shapes..csrrqrrr?r?r@ru*rr)rFZget_input_shapeZget_padded_shaperprrrQ) r>r4rroriginal_spatial_shapeZpadded_spatial_shapeoutput_spatial_shapeZpadded_output_spatial_shaperrr?r?r@rs   zPatchInferer._get_merged_shapesr4NCallable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]]cOs~|dd}|durpt|tjr)t|tjr)|j|jkr(td|jd|jnGt|trlt|trlt|t|krJtdt|dt|dt||D]\\}}\}}|j|jkrjtd|jd|jdqOntd |j durt|tjrt|t rt j |j vrtd n td t|d|} |dur|} n| |} |dur| |} g} g} |durt|| || D]0\\} }}\}}}||d<|j|| g|Ri|}| s|||| |\} } ||||| | qn.|| D](\} }}|j|| g|Ri|}| s|||| |\} } ||||| | qd d | D}|jr2tt|j|St|dkr=|dS|S)a Args: inputs: input data for inference, a torch.Tensor, representing an image or batch of images. However if the data is already split, it can be fed by providing a list of tuple (patch, location), or a MetaTensor that has metadata for `PatchKeys.LOCATION`. In both cases no splitter should be provided. network: target model to execute inference. supports callables such as ``lambda x: my_torch_model(x, additional_config)`` args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. condition (torch.Tensor, optional): If provided via `**kwargs`, this tensor must match the shape of `inputs` and will be sliced, patched, or windowed alongside the inputs. The resulting segments will be passed to the model together with the corresponding input segments. conditionN*`condition` must match shape of `inputs` ( ), but got z/Length of `condition` must match `inputs`. Got  and .zREach `condition` patch must match the shape of the corresponding input patch. Got zX`condition` and `inputs` must be of the same type (both Tensor or both list of patches).z`PatchKey.LOCATION` does not exists in `inputs.meta`. If the inputs are already split into patches, the location of patches needs to be provided as `PatchKey.LOCATION` metadata in a MetaTensor. If the input is not already split, please provide `splitter`.z`splitter` should be set if the input is not already split into patches. For inputs that are split, the location of patches needs to be provided as (image, location) pairs, or as `PatchKey.LOCATION` metadata in a MetaTensor. The provided inputs type is cSsg|]}|qSr?)finalize)rrrr?r?r@rz)PatchInferer.__call__..rEr)poprWrgTensorrr[rxrbrrFrr#rfrerXrnr|rrrOrw)r>r4r6r7r8rr_Z cond_patchZpatches_locationsZcondition_locationsrrr^rrJZcondition_patchesroZmerged_outputsr?r?r@rA5s          zPatchInferer.__call__)rFrGrHrIrJrKrLrMrNrMrOrPrQrRrSrKrTr r9rU)r^r_r9r`)ror r9rp) r6rr{r5r7r r8r r9rp) r4r5r6rr7r r8r r9r ) r<rBrCrDrrVrnrzr|rrrrAr?r?r?r@r.ds$ ? (  r.c@s$eZdZdZdddZdd dZdS)r/z SimpleInferer is the normal inference method that run model forward() directly. Usage example can be found in the :py:class:`monai.inferers.Inferer` base class. r9rUcCt|dSrqr-rVr>r?r?r@rVzSimpleInferer.__init__r4r5r6Callable[..., torch.Tensor]r7r r8cOs||g|Ri|S)aUnified callable function API of Inferers. Args: inputs: model input data for inference. network: target model to execute inference. supports callables such as ``lambda x: my_torch_model(x, additional_config)`` args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. r?r=r?r?r@rAs zSimpleInferer.__call__Nr9rU) r4r5r6rr7r r8r r9r5r<rBrCrDrVrAr?r?r?r@r/s r/csNeZdZdZddejdejdddddddddfd/fd#d$ Zd0d-d.ZZ S)1r0a Sliding window method for model inference, with `sw_batch_size` windows for every model.forward(). Usage example can be found in the :py:class:`monai.inferers.Inferer` base class. Args: roi_size: the window size to execute SlidingWindow evaluation. If it has non-positive components, the corresponding `inputs` size will be used. if the components of the `roi_size` are non-positive values, the transform will use the corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted to `(32, 64)` if the second spatial dimension size of img is `64`. sw_batch_size: the batch size to run window slices. overlap: Amount of overlap between scans along each spatial dimension, defaults to ``0.25``. mode: {``"constant"``, ``"gaussian"``} How to blend output of overlapping windows. Defaults to ``"constant"``. - ``"constant``": gives equal weight to all predictions. - ``"gaussian``": gives less weight to predictions on edges of windows. sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``. Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``. When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding spatial dimensions. padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``} Padding mode when ``roi_size`` is larger than inputs. Defaults to ``"constant"`` See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html cval: fill value for 'constant' padding mode. Default: 0 sw_device: device for the window data. By default the device (and accordingly the memory) of the `inputs` is used. Normally `sw_device` should be consistent with the device where `predictor` is defined. device: device for the stitched output prediction. By default the device (and accordingly the memory) of the `inputs` is used. If for example set to device=torch.device('cpu') the gpu memory consumption is less and independent of the `inputs` and `roi_size`. Output is on the `device`. progress: whether to print a tqdm progress bar. cache_roi_weight_map: whether to precompute the ROI weight map. cpu_thresh: when provided, dynamically switch to stitching on cpu (to save gpu memory) when input image volume is larger than this threshold (in pixels/voxels). Otherwise use ``"device"``. Thus, the output may end-up on either cpu or gpu. buffer_steps: the number of sliding window iterations along the ``buffer_dim`` to be buffered on ``sw_device`` before writing to ``device``. (Typically, ``sw_device`` is ``cuda`` and ``device`` is ``cpu``.) default is None, no buffering. For the buffer dim, when spatial size is divisible by buffer_steps*roi_size, (i.e. no overlapping among the buffers) non_blocking copy may be automatically enabled for efficiency. buffer_dim: the spatial dimension along which the buffers are created. 0 indicates the first spatial dimension. Default is -1, the last spatial dimension. with_coord: whether to pass the window coordinates to ``network``. Defaults to False. If True, the ``network``'s 2nd input argument should accept the window coordinates. Note: ``sw_batch_size`` denotes the max number of windows per network inference iteration, not the batch size of inputs. rEg?g?gNFroi_sizeSequence[int] | int sw_batch_sizerKoverlapSequence[float] | floatmodeBlendMode | str sigma_scale padding_modePytorchPadMode | strcvalfloat sw_devicetorch.device | str | NonedeviceprogressrRcache_roi_weight_map cpu_thresh int | None buffer_steps buffer_dim with_coordr9rUc st||_||_||_t||_||_||_||_ ||_ | |_ | |_ | |_ | |_||_||_d|_z4| rTt|trTt|dkrT| durHd} tt|j||| d|_| rc|jdurftdWdSWdSWdSty}ztd|jd|d|d| d |d}~ww) Nrcpu)rrrzHcache_roi_weight_map=True, but cache is not created. (dynamic roi_size?)z roi size z, mode=z, sigma_scale=z , device=z^ Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'.)superrVrrrr!rrrrrrrrrrrroi_weight_maprWr rdrr%warningswarn BaseException RuntimeError)r>rrrrrrrrrrrrrrrer;r?r@rVsD   zSlidingWindowInferer.__init__r4r5r6rr7r r8Atorch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]c Os|dd}|dur|j|jkrtd|jd|j|d|j}|d|j}|d|j}|durH|jdurH|jdd|jkrHd }t ||j |j ||j |j |j|j|j|j||j|jd|||jg|Ri|S) a Args: inputs: model input data for inference. network: target model to execute inference. supports callables such as ``lambda x: my_torch_model(x, additional_config)`` args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. condition (torch.Tensor, optional): If provided via `**kwargs`, this tensor must match the shape of `inputs` and will be sliced, patched, or windowed alongside the inputs. The resulting segments will be passed to the model together with the corresponding input segments. rNrrrrrrr)getrr[rrrrrnumelrrrrrrrrrrrr) r>r4r6r7r8rrrrr?r?r@rA*s> *zSlidingWindowInferer.__call__) rrrrKrrrrrrrrrrrrrrrrRrrRrrrrrrKrrRr9rU r4r5r6rr7r r8r r9r) r<rBrCrDr!CONSTANTr$rVrA __classcell__r?r?rr@r0s$:4r0cs"eZdZdZd fd d ZZS)r3a( SlidingWindowInfererAdapt extends SlidingWindowInferer to automatically switch to buffered and then to CPU stitching, when OOM on GPU. It also records a size of such large images to automatically try CPU stitching for the next large image of a similar size. If the stitching 'device' input parameter is provided, automatic adaptation won't be attempted, please keep the default option device = None for adaptive behavior. Note: the output might be on CPU (even if the input was on GPU), if the GPU memory was not sufficient. r4r5r6rr7r r8r9rc sV|jdurtj||g|Ri|S|jduo|jdk}|jduo-|jdd|jk}|jo3| }|jo;|o;| }|jdurGtd|jnd} d} t |jdd} | t| } |j| d|jddkrk| } t dD]} z!tj||g|R|r|jnt d|r| nd| d|WSt y}zx|s|rd tt|jvr|t||rd }|jddd|_|rd }td |jd n?d }| |_td| d| d|jd n)| dkrtd| d} | |_td|jd| d n d }td|jd WYd}~qod}~wwt d|d|d|d|d| )ag Args: inputs: model input data for inference. network: target model to execute inference. supports callables such as ``lambda x: my_torch_model(x, additional_config)`` args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. NrrrEr r)rrrOutOfMemoryErrorFz3GPU stitching failed, attempting on CPU, image dim rTzGPU stitching failed, buffer z dim z , image dim z)GPU buffered stitching failed, image dim z reducing buffer to zr4r6r7r8Z skip_bufferZcpu_condZ gpu_stitchingZbuffered_stitchingrrshmax_dimrrrr?r@rAjsl "        z"SlidingWindowInfererAdapt.__call__r)r<rBrCrDrArr?r?rr@r3`s r3c@s(eZdZdZ ddd dZdddZdS)r1a SaliencyInferer is inference with activation maps. Args: cam_name: expected CAM method name, should be: "CAM", "GradCAM" or "GradCAMpp". target_layers: name of the model layer to generate the feature map. class_idx: index of the class to be visualized. if None, default to argmax(logits). args: other optional args to be passed to the `__init__` of cam. kwargs: other optional keyword args to be passed to `__init__` of cam. Ncam_namerZ target_layers class_idxrr7r r8r9rUcOsDt||dvrtd||_||_||_||_||_dS)N)camgradcamZ gradcamppz4cam_name should be: 'CAM', 'GradCAM' or 'GradCAMpp'.) r-rVlowerr[rrrr7r8)r>rrrr7r8r?r?r@rVs    zSaliencyInferer.__init__r4r5r6 nn.ModulecOs|jdkrt||jg|jRi|j}n$|jdkr*t||jg|jRi|j}nt||jg|jRi|j}|||jg|Ri|S)aUnified callable function API of Inferers. Args: inputs: model input data for inference. network: target model to execute inference. supports callables such as ``lambda x: my_torch_model(x, additional_config)`` args: other optional args to be passed to the `__call__` of cam. kwargs: other optional keyword args to be passed to `__call__` of cam. rr)rr'rr7r8r(r)r)r>r4r6r7r8rr?r?r@rAs   zSaliencyInferer.__call__rq) rrZrrZrrr7r r8r r9rU)r4r5r6rr7r r8r rr?r?r?r@r1s   r1cs@eZdZdZddfd d Zdfdd Z ddddZZS)r2a[ SliceInferer extends SlidingWindowInferer to provide slice-by-slice (2D) inference when provided a 3D volume. A typical use case could be a 2D model (like 2D segmentation UNet) operates on the slices from a 3D volume, and the output is a 3D volume with 2D slices aggregated. Example:: # sliding over the `spatial_dim` inferer = SliceInferer(roi_size=(64, 256), sw_batch_size=1, spatial_dim=1) output = inferer(input_volume, net) Args: spatial_dim: Spatial dimension over which the slice-by-slice inference runs on the 3D volume. For example ``0`` could slide over axial slices. ``1`` over coronal slices and ``2`` over sagittal slices. args: other optional args to be passed to the `__init__` of base class SlidingWindowInferer. kwargs: other optional keyword args to be passed to `__init__` of base class SlidingWindowInferer. Note: ``roi_size`` in SliceInferer is expected to be a 2D tuple when a 3D volume is provided. This allows sliding across slices along the 3D volume using a selected ``spatial_dim``. r spatial_dimrKr7r r8r9rUcs(||_tj|i|t|j|_dSrq)rrrVr%r orig_roi_size)r>rr7r8rr?r@rVszSliceInferer.__init__r4r5r6rrcsjdkr tdtj_tjdkr0t|jdddkr0tj_jjdn t djd|jd| d d}|durY|j|jkrYtd |jd |j|durkt j |fd d |dSt j |fdd dS)aE Args: inputs: 3D input for inference network: 2D model to execute inference on slices in the 3D input args: optional args to be passed to ``network``. kwargs: optional keyword args to be passed to ``network``. condition (torch.Tensor, optional): If provided via `**kwargs`, this tensor must match the shape of `inputs` and will be sliced, patched, or windowed alongside the inputs. The resulting segments will be passed to the model together with the corresponding input segments.rzB`spatial_dim` can only be `0, 1, 2` with `[H, W, D]` respectively.NrEzCurrently, only 2D `roi_size` (z!) with 3D `inputs` tensor (shape=z) is supported.rrrcj|g|Ri|Srqnetwork_wrapperxr7r8r6r>r?r@)z'SliceInferer.__call__..)r4r6rcrrqrrrr?r@r.r)r4r6) rr[r%rrbrrrxinsertrrrrA)r>r4r6r7r8rrrr@rAs*  $   zSliceInferer.__call__Nrrtorch.Tensor | Nonecs|jjdd}|dur#|jjdd}|||g|Ri|}n ||g|Ri|}t|tjr=|jjddSt|trX|D]}||jjdd||<qF|Stfdd|DS)zP Wrapper handles inference for 2D models over 3D volume inputs. rdimNc3s"|] }|jjddVqdS)rrN) unsqueezer)rrout_irr?r@ruOs z/SliceInferer.network_wrapper..) squeezerrWrgrrrryrp)r>r6rrr7r8outrsr?rr@r1s    zSliceInferer.network_wrapper)r)rrKr7r r8r r9rUrrq) r6rrr5rrr7r r8r r9r)r<rBrCrDrVrArrr?r?rr@r2s 2r2cseZdZdZd8fdd Z  d9d:ddZe        d;dd0d1Z d2d3Z + ,d?d@d6d7Z Z S)ADiffusionInfererz DiffusionInferer takes a trained diffusion model and a scheduler and can be used to perform a signal forward pass for a training iteration, and sample from the model. Args: scheduler: diffusion scheduler. schedulerrr9rUcst||_dSrq)rrVrr>rrr?r@rV[s  zDiffusionInferer.__init__N crossattnr4r5diffusion_modelrnoise timestepsrrrrZsegc Cs|dvr t|d|jj|||d}|dkr+|dur tdtj||gdd}d}t|tr6t||d n|}||||d } | S) a> Implements the forward pass for a supervised training iteration. Args: inputs: Input image to which noise is added. diffusion_model: diffusion model. noise: random noise, of the same shape as the input. timesteps: random timesteps. condition: Conditioning for network input. mode: Conditioning mode for the network. seg: if model is instance of SPADEDiffusionModelUnet, segmentation must be provided on the forward (for SPADE-like AE or SPADE-like DM) rconcat condition is not supportedZoriginal_samplesrrrNz-Conditioning is required for concat conditionrErrrrcontext) r:r add_noiser[rgrhrWrr ) r>r4rrrrrr noisy_image predictionr?r?r@rA`szDiffusionInferer.__call__FdT input_noiseScheduler | Nonesave_intermediates bool | Noneintermediate_stepsr conditioningverboserRcfg float | Nonecfg_fill_valuer6torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]c Cs |dvr t|d|dkr|durtd|s|j}|} t|jddtjdg|jjdf} |rItrIt t |j| t t |jt | d }nt t |j| }g}|D]\}}t|trdt|| d n|}| durtj| gd dd }|durt|}|| tj||gdd }nd}n| }|}|dkr|durtj||gdd }||t|f|jdd }n||t|f|j|d }| dur|d \}}|| ||}t|ts|||| \} }n |||| |\} }|r||dkr|| qU|r| |fS| S)a Args: input_noise: random noise, of the same shape as the desired sample. diffusion_model: model to sample from. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler save_intermediates: whether to return intermediates along the sampling change intermediate_steps: if save_intermediates is True, saves every n steps conditioning: Conditioning for network input. mode: Conditioning mode for the network. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. rrrN>Conditioning must be supplied for if condition mode is concat.rErdtypetotalrrrrr)r:r[rrgrhrtensorrhas_tqdmr*rrdrbiterrWrr ones_likefill_rtorrrstepr)r>rrrrrrrrrrrimageall_next_timesteps progress_bar intermediatestnext_t model_input unconditionconditioning_input model_outputmodel_output_uncondmodel_output_condrr?r?r@rmsb(      zDiffusionInferer.samplerrrEoriginal_input_rangerpscaled_input_rangec !Cs8|s|j}|dkrtd||dvrt|d|dkr+|dur+td| r5tr5t|j} nt|j} g} t | |j } t |j d |j }| D]>}tj|j dd ||j d }|jj|| |d }t|trzt|| d n|}|dkr|durtj||gd d }|||dd}n||||d}|j d |j d dkr|jdvrtj||j d d d \}}nd}|j|}|dkr|j|d n|j}d |}d |}|jdkr||d||d}n|jdkr|}n|jdkr|d||d|}|jrt|dd }|d|j||}|j|d||}||||}|j|||d}|j||d}t |}|rIt |n|}|dkr^|j!||d|||d } ndd||t"||||dt"| } || #| j ddj$d d 7}|r| %| &qS|r|| fS|S)a[ Computes the log-likelihoods for an input. Args: inputs: input images, NxCxHxW[xD] diffusion_model: model to compute likelihood from scheduler: diffusion scheduler. If none provided will use the class attribute scheduler. save_intermediates: save the intermediate spatial KL maps conditioning: Conditioning for network input. mode: Conditioning mode for the network. original_input_range: the [min,max] intensity range of the input data before any scaling was applied. scaled_input_range: the [min,max] intensity range of the input data after scaling. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. DDPMSchedulerLLikelihood computation is only compatible with DDPMScheduler, you are using rrrNrrrErrrrrrrZlearnedZ learned_rangeepsilon?rm v_predictionrtimestepx_0x_tr?predicted_variancer4means log_scalesr5r6r )'r _get_namer:r[r r*rr!rg randn_liker$rzerosrfulllongr rWrr rh variance_typesplitalphas_cumprodoneprediction_type clip_sampleclampbetasalphas _get_mean _get_variancelog_get_decoder_log_likelihoodexpviewmeanrr)!r>r4rrrrrr5r6rrr(r)rtotal_klr*rr r/rC alpha_prod_talpha_prod_t_prev beta_prod_tbeta_prod_t_prevpred_original_samplepred_original_sample_coeffcurrent_sample_coeffpredicted_meanposterior_meanposterior_variancelog_posterior_variancelog_predicted_varianceklr?r?r@get_likelihoods     "         zDiffusionInferer.get_likelihoodc CsBddtttdtjg|j|dt|dS)z A fast approximation of the cumulative distribution function of the standard normal. Code adapted from https://github.com/openai/improved-diffusion. r<?g@gHm?r) rgtanhsqrtrmathpir$rpow)r>rr?r?r@_approx_standard_normal_cdfhs<z,DiffusionInferer._approx_standard_normal_cdfrErFc Cs|j|jkrtd|jd|j|d|d|d|d}||}t| }|||d} || } |||d} || } t| jdd} td| jdd}| | }t|d k| t|d k|t|jdd}|S) ax Compute the log-likelihood of a Gaussian distribution discretizing to a given image. Code adapted from https://github.com/openai/improved-diffusion. Args: input: the target images. It is assumed that this was uint8 values, rescaled to the range [-1, 1]. means: the Gaussian mean Tensor. log_scales: the Gaussian log stddev Tensor. original_input_range: the [min,max] intensity range of the input data before any scaling was applied. scaled_input_range: the [min,max] intensity range of the input data after scaling. z/Inputs and means must have the same shape, got rrErrg-q=)rdrkg+g+?)rr[rgrYrqrWrRwhere)r>r4rErFr5r6 bin_widthZ centered_xZinv_stdvZplus_inZcdf_plusZmin_inZcdf_minZ log_cdf_plusZlog_one_minus_cdf_minZ cdf_delta log_probsr?r?r@rXrs(    z,DiffusionInferer._get_decoder_log_likelihoodrrr9rUNrN)r4r5rrrr5rr5rrrrZrrr9r5 NFr NrTNNr )rr5rrrrrrrrrrrrZrrRrrrrrrr9rNFNrr2r4TN)r4r5rrrrrrrrrrZr5rpr6rprrRrrr9r)r2r4) r4r5rEr5rFr5r5rpr6rpr9r5) r<rBrCrDrVrArgno_gradrmrjrqrXrr?r?rr@rRsB  *\rcseZdZdZ   dd?fdd Ze     !   "d@dAfd0d1 Ze     2 3 !  4 dBdCfd:d; Z Z S)DLatentDiffusionInferera- LatentDiffusionInferer takes a stage 1 model (VQVAE or AutoencoderKL), diffusion model, and a scheduler, and can be used to perform a signal forward pass for a training iteration, and sample from the model. Args: scheduler: a scheduler to be used in combination with `unet` to denoise the encoded image latents. scale_factor: scale factor to multiply the values of the latent representation before processing it by the second stage. ldm_latent_shape: desired spatial latent space shape. Used if there is a difference in the autoencoder model's latent shape. autoencoder_latent_shape: autoencoder_latent_shape: autoencoder spatial latent space shape. Used if there is a difference between the autoencoder's latent shape and the DM shape. rkNrr scale_factorrldm_latent_shape list | Noneautoencoder_latent_shaper9rUcttj|d||_|du|duArtd||_||_|jdur6|jdur8t|jd|_t|jd|_ dSdSdS)NrzSIf ldm_latent_shape is None, autoencoder_latent_shape must be None, and vice versa. spatial_sizer rrVr{r[r|r~r ldm_resizerrautoencoder_resizerr>rr{r|r~rr?r@rVzLatentDiffusionInferer.__init__rr4r5autoencoder_modelAutoencoderKL | VQVAErrrrrrrrZrc s~t||j} Wdn1swYjdur0tfddt| Dd} tj| ||||||d} | S)aw Implements the forward pass for a supervised training iteration. Args: inputs: input image to which the latent representation will be extracted and noise is added. autoencoder_model: first stage model. diffusion_model: diffusion model. noise: random noise, of the same shape as the latent representation. timesteps: random timesteps. condition: conditioning for network input. mode: Conditioning mode for the network. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. Ncg|]}|qSr?rrrrjrr?r@rz3LatentDiffusionInferer.__call__..r)r4rrrrrr) rgryencode_stage_2_inputsr{r|stackrrrA) r>r4rrrrrrrlatentr rrr@rAs   zLatentDiffusionInferer.__call__Fr Tr rrrrrrrrrRrrrrc  s*t|trt|tr|jj|jkrtd|jjd|jtj|||||||| | | | d } |r6| \}}n| }jdurWt fddt |Dd}|rWfdd|D}|j }t|trft |j | d }||j}|rg}|D]}|j }t|trt |j | d }|||jqs||fS|S) a> Args: input_noise: random noise, of the same shape as the desired latent representation. autoencoder_model: first stage model. diffusion_model: model to sample from. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler. save_intermediates: whether to return intermediates along the sampling change intermediate_steps: if save_intermediates is True, saves every n steps conditioning: Conditioning for network input. mode: Conditioning mode for the network. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, or autoencoder_model is instance of SPADEAutoencoderKL, segmentation must be provided. cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. zIf both autoencoder_model and diffusion_model implement SPADE, the number of semanticlabels for each must be compatible, but got z and) rrrrrrrrrrrNcrr?rrrr?r@r$rz1LatentDiffusionInferer.sample..rc*g|]}tfddt|DdqS)crr?rrrr?r@r'rz..rrgrrrrrrr?r@r&r)rWrrdecoderlabel_ncr[rrmr~rgrrdecode_stage_2_outputsr r{r)r>rrrrrrrrrrrrrorlatent_intermediatesdecoder&r)latent_intermediaterrr@rms^"     zLatentDiffusionInferer.sampler2r4nearestr5 tuple | Noner6resample_latent_likelihoodsresample_interpolation_modec s| r | dvr td| ||j}jdur)tfddt|Dd}tj||||||| | d}|rZ| rZ|d}t j |j d d| d fd d|D}|d|f}|S) a, Computes the log-likelihoods of the latent representations of the input. Args: inputs: input images, NxCxHxW[xD] autoencoder_model: first stage model. diffusion_model: model to compute likelihood from scheduler: diffusion scheduler. If none provided will use the class attribute scheduler save_intermediates: save the intermediate spatial KL maps conditioning: Conditioning for network input. mode: Conditioning mode for the network. original_input_range: the [min,max] intensity range of the input data before any scaling was applied. scaled_input_range: the [min,max] intensity range of the input data after scaling. verbose: if true, prints the progression bar of the sampling process. resample_latent_likelihoods: if true, resamples the intermediate likelihood maps to have the same spatial dimension as the input images. resample_interpolation_mode: if use resample_latent_likelihoods, select interpolation 'nearest', 'bilinear', or 'trilinear; seg: if diffusion model is instance of SPADEDiffusionModel, or autoencoder_model is instance of SPADEAutoencoderKL, segmentation must be provided. rbilinear trilinearRresample_interpolation mode should be either nearest, bilinear, or trilinear, got Ncrr?rrrr?r@rhrz9LatentDiffusionInferer.get_likelihood..r)r4rrrrrrrrErsizercg|]}|qSr?r?rrrresizerr?r@rxr) r[rr{r|rgrrrrjnnUpsampler)r>r4rrrrrrr5r6rrrrlatentsror)rrr>r@rj;s. &   z%LatentDiffusionInferer.get_likelihoodrkNN rrr{rr|r}r~r}r9rUrv)r4r5rrrrrr5rr5rrrrZrrr9r5rw)rr5rrrrrrrrrrrrrrZrrRrrrrrrr9r NFNrr2r4TFrN)r4r5rrrrrrrrrrrrZr5rr6rrrRrrRrrZrrr9r r<rBrCrDrVrArgryrmrjrr?r?rr@rzsB)Vrzcs~eZdZdZd5ddZ  d6d7ddZe        d8d9d,d-Ze    . /  d:d;fd3d4 Z Z S)<ControlNetDiffusionInferera ControlNetDiffusionInferer takes a trained diffusion model and a scheduler and can be used to perform a signal forward pass for a training iteration, and sample from the model, supporting ControlNet-based conditioning. Args: scheduler: diffusion scheduler. rrr9rUcCst|||_dSrq)r-rVrrr?r?r@rVs  z#ControlNetDiffusionInferer.__init__Nrr4r5rr controlnetrrrcn_condrrrrZrc Cs|dvr t|d|jj|||d} |dkr'|dur'tj| |gdd} d}|| |||d\} } |} t|tr>t|| d } | | ||| | d }|S) a Implements the forward pass for a supervised training iteration. Args: inputs: Input image to which noise is added. diffusion_model: diffusion model. controlnet: controlnet sub-network. noise: random noise, of the same shape as the input. timesteps: random timesteps. cn_cond: conditioning image for the ControlNet. condition: Conditioning for network input. mode: Conditioning mode for the network. seg: if model is instance of SPADEDiffusionModelUnet, segmentation must be provided on the forward (for SPADE-like AE or SPADE-like DM) rrrrNrErrrcontrolnet_condrrrrrdown_block_additional_residualsmid_block_additional_residual)r:rr rgrhrWrr )r>r4rrrrrrrrr down_block_res_samplesmid_block_res_samplediffuser r?r?r@rAs(   z#ControlNetDiffusionInferer.__call__Fr Tr rrrrrrrrrRrrrrrcCsh| dvr t| d|s|j}|}t|jddtjdg|jjdf}| r=tr=tt |j|t t |jt |d}nt t |j|}g}| durUtj|gddd }|D]\}}| durtj|gddd }|durt |}|| tj||gdd }nd}n|}|}|}t|trt|| d }| d kr|durtj||gdd }||t|f|j|dd \}}||t|f|jd||d }n#||t|f|j||d \}}||t|f|j|||d }| dur|d\}}|| ||}t|ts||||\}}n |||||\}}|r*||dkr*||qW|r2||fS|S)a# Args: input_noise: random noise, of the same shape as the desired sample. diffusion_model: model to sample from. controlnet: controlnet sub-network. cn_cond: conditioning image for the ControlNet. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler save_intermediates: whether to return intermediates along the sampling change intermediate_steps: if save_intermediates is True, saves every n steps conditioning: Conditioning for network input. mode: Conditioning mode for the network. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. rrrENrrrrrrrrrrrr)r:rrgrhrrrr r*rrdrbr!r"r#rWrr rr$rrrr%r)r>rrrrrrrrrrrrrr&r'r(r)r*r+r,r-r.rrrr/r0r1rr?r?r@rms!(            z!ControlNetDiffusionInferer.sampler2r4r5rpr6c &sv|s|j}|dkrtd||dvrt|d| r)tr)t|j} nt|j} g}t| |j }t |j d |j }| D]i}tj |j dd||j d}|jj|||d }|}t|trpt|| d }|d kr|durtj||gdd }||t|f |j |dd \}}|||d||d}n||t|f |j ||d \}}||||||d}|j d|j ddkr|jdvrtj||j ddd \}}nd}|j|}|dkr|j|dn|j}d|}d|}|jdkr||d||d}n|jdkr|}n|jdkr"|d||d|}|jr-t|dd}|d|j||}|j|d||}||||} |j|||d}!|j||d}"t |"}#|rgt |n|#}$|dkr}t!j"|| d|$| | d }%ndd|$|#t#|#|$|!| dt#|$ }%||%$|%j ddj%dd 7}|r|&|%'qG|r||fS|S)a Computes the log-likelihoods for an input. Args: inputs: input images, NxCxHxW[xD] diffusion_model: model to compute likelihood from controlnet: controlnet sub-network. cn_cond: conditioning image for the ControlNet. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler. save_intermediates: save the intermediate spatial KL maps conditioning: Conditioning for network input. mode: Conditioning mode for the network. original_input_range: the [min,max] intensity range of the input data before any scaling was applied. scaled_input_range: the [min,max] intensity range of the input data after scaling. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. r7r8rrrNrEr9rrrrrrrrr:r;r<rmr=rr>rBrDr )(rrGr:r r*rr!rgrHr$rrIrrJrKr rWrr rhrrLrMrNrOrPrQrRrSrTrUrVrWrrXrYrZr[rr)&r>r4rrrrrrrr5r6rrr(r)rr\r*rr rrrr/rCr]r^r_r`rarbrcrdrerfrgrhrirr?r@rj>s"         "         z)ControlNetDiffusionInferer.get_likelihoodrurv)r4r5rrrrrr5rr5rr5rrrrZrrr9r5rw)rr5rrrrrr5rrrrrrrrrrZrrRrrrrrrr9rrx)r4r5rrrrrr5rrrrrrrrZr5rpr6rprrRrrr9rrr?r?rr@r}s8   6}rcseZdZdZ   d?d@fd d Z   dAdBfd d! Ze  " #   $   %dCdDfd3d4 Ze  "   5 6 $ " 7 dEdFfd=d> Z Z S)G ControlNetLatentDiffusionInfereraG ControlNetLatentDiffusionInferer takes a stage 1 model (VQVAE or AutoencoderKL), diffusion model, controlnet, and a scheduler, and can be used to perform a signal forward pass for a training iteration, and sample from the model. Args: scheduler: a scheduler to be used in combination with `unet` to denoise the encoded image latents. scale_factor: scale factor to multiply the values of the latent representation before processing it by the second stage. ldm_latent_shape: desired spatial latent space shape. Used if there is a difference in the autoencoder model's latent shape. autoencoder_latent_shape: autoencoder_latent_shape: autoencoder spatial latent space shape. Used if there is a difference between the autoencoder's latent shape and the DM shape. rkNrrr{rr|r}r~r9rUcr)NrzQIf ldm_latent_shape is None, autoencoder_latent_shape must be Noneand vice versa.rrrrrr?r@rVrz)ControlNetLatentDiffusionInferer.__init__rr4r5rrrrrrrrrrrrrZrc st||j} Wdn1swYjdur0tfddt| Dd} |jdd| jddkrIt || jdd}t j | ||||||| | d } | S)a Implements the forward pass for a supervised training iteration. Args: inputs: input image to which the latent representation will be extracted and noise is added. autoencoder_model: first stage model. diffusion_model: diffusion model. controlnet: instance of ControlNet model noise: random noise, of the same shape as the latent representation. timesteps: random timesteps. cn_cond: conditioning tensor for the ControlNet network condition: conditioning for network input. mode: Conditioning mode for the network. seg: if diffusion model is instance of SPADEDiffusionModel, segmentation must be provided. Ncrr?rrrr?r@rrz=ControlNetLatentDiffusionInferer.__call__..rr) r4rrrrrrrr) rgryrr{r|rrrF interpolaterrA) r>r4rrrrrrrrrrr rrr@rAs&   z)ControlNetLatentDiffusionInferer.__call__Fr Tr rrrrrrrrrRrrrrcsNt|trt|tr|jj|jkrtd|jdd|jddkr.t||jdd}t j |||||||| | | | | |d }|rH|\}}n|}j durit fddt|Dd}|rifdd|D}|j}t|trxt|j| d }||j}|rg}|D]}|j}t|trt|j| d }|||jq||fS|S) a Args: input_noise: random noise, of the same shape as the desired latent representation. autoencoder_model: first stage model. diffusion_model: model to sample from. controlnet: instance of ControlNet model. cn_cond: conditioning tensor for the ControlNet network. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler. save_intermediates: whether to return intermediates along the sampling change intermediate_steps: if save_intermediates is True, saves every n steps conditioning: Conditioning for network input. mode: Conditioning mode for the network. verbose: if true, prints the progression bar of the sampling process. seg: if diffusion model is instance of SPADEDiffusionModel, or autoencoder_model is instance of SPADEAutoencoderKL, segmentation must be provided. cfg: classifier-free-guidance scale, which indicates the level of strengthening on the conditioning. cfg_fill_value: the fill value to use for the unconditioned input when using classifier-free guidance. zIf both autoencoder_model and diffusion_model implement SPADE, the number of semanticlabels for each must be compatible. Got {autoencoder_model.decoder.label_nc} and {diffusion_model.label_nc}rN) rrrrrrrrrrrrrcrr?rrrr?r@rqrz;ControlNetLatentDiffusionInferer.sample..rcr)crr?rrrr?r@rtrzFControlNetLatentDiffusionInferer.sample...rrrrr?r@rsrr)rWrrrrr[rrrrrmr~rgrrrr r{r)r>rrrrrrrrrrrrrrrorrrr&r)rrrr@rm)s^&     z'ControlNetLatentDiffusionInferer.sampler2r4rr5rr6rrc s| r |dvr td|||j}|jdd|jddkr.t||jdd}jdurBtfddt |Dd}t j |||||||| | |d }|ru| ru|d }t j |jdd|d fd d|D}|d|f}|S) a Computes the log-likelihoods of the latent representations of the input. Args: inputs: input images, NxCxHxW[xD] autoencoder_model: first stage model. diffusion_model: model to compute likelihood from controlnet: instance of ControlNet model. cn_cond: conditioning tensor for the ControlNet network. scheduler: diffusion scheduler. If none provided will use the class attribute scheduler save_intermediates: save the intermediate spatial KL maps conditioning: Conditioning for network input. mode: Conditioning mode for the network. original_input_range: the [min,max] intensity range of the input data before any scaling was applied. scaled_input_range: the [min,max] intensity range of the input data after scaling. verbose: if true, prints the progression bar of the sampling process. resample_latent_likelihoods: if true, resamples the intermediate likelihood maps to have the same spatial dimension as the input images. resample_interpolation_mode: if use resample_latent_likelihoods, select interpolation 'nearest', 'bilinear', or 'trilinear; seg: if diffusion model is instance of SPADEDiffusionModel, or autoencoder_model is instance of SPADEAutoencoderKL, segmentation must be provided. rrrNcrr?rrrr?r@rrzCControlNetLatentDiffusionInferer.get_likelihood..r) r4rrrrrrrrrrErcrr?r?rrr?r@rr)r[rr{rrrr|rgrrrrjrr)r>r4rrrrrrrrr5r6rrrrrror)rrr@rjs6 *   z/ControlNetLatentDiffusionInferer.get_likelihoodrrrv)r4r5rrrrrrrr5rr5rr5rrrrZrrr9r5rw)rr5rrrrrrrr5rrrrrrrrrrZrrRrrrrrrr9rr) r4r5rrrrrrrr5rrrrrrrrZr5rr6rrrRrrRrrZrrr9rrr?r?rr@rsB3`rc@sbeZdZdZd*ddZ  d+d,ddZe    d-d.d"d#Ze   $ d/d0d(d)Z dS)1VQVAETransformerInfererzF Class to perform inference with a VQVAE + Transformer model. r9rUcCrrqrrr?r?r@rVrz VQVAETransformerInferer.__init__NFr4r5 vqvae_modelrtransformer_modelrorderingr"rr return_latentrR7torch.Tensor | tuple[torch.Tensor, torch.Tensor, tuple]cCs$t ||}Wdn1swYt|jdd}||jdd}|dd|f}|} t |dd|j }|ddddf}| }|jd} |j } | | krmt tjd| d| dd} nd} ||dd| | | f|d } |r| | dd| | | f|fS| S) a Implements the forward pass for a supervised training iteration. Args: inputs: input image to which the latent representation will be extracted. vqvae_model: first stage model. transformer_model: autoregressive transformer model. ordering: ordering of the quantised latent representation. return_latent: also return latent sequence and spatial dim of the latent. condition: conditioning for network input. NrErrrErconstant)rE)lowhighrrr)rgryindex_quantizerprreshapeget_sequence_orderingclonerpadnum_embeddingsrK max_seq_lenrKrandintitem)r>r4rrrrrrlatent_spatial_dimtargetseq_lenrstartr r?r?r@rAs&   " z VQVAETransformerInferer.__call__rkTr&tuple[int, int, int] | tuple[int, int]starting_tokensr temperaturertop_krrc  Cs^t|} | rtrtt| } ntt| } |} | D]n} | d|jkr)| }n | dd|j df}|||d}|dddddf|}|durjt |t ||d\}} t d |||dddgfk<t j|dd}d|dd|jf<t j|dd}t j| |fdd} q| ddddf} | dd|f} | |jdf|}||S) a@ Sampling function for the VQVAE + Transformer model. Args: latent_spatial_dim: shape of the sampled image. starting_tokens: starting tokens for the sampling. It must be vqvae_model.num_embeddings value. vqvae_model: first stage model. transformer_model: model to sample from. conditioning: Conditioning for network input. temperature: temperature for sampling. top_k: top k sampling. verbose: if true, prints the progression bar of the sampling process. rENrrInfrr) num_samples)rnprodr r*rcr!rKrrrgtopkrdrrsoftmaxr multinomialrhget_revert_sequence_orderingrrdecode_samples)r>rrrrrrrrrrr(Z latent_seqridx_condlogitsvprobsZidx_nextrr?r?r@rms,     zVQVAETransformerInferer.samplerrrrZc CsL|r |dvr td|t ||} Wdn1s!wYt| jdd} | | jdd} | dd|f} t | } t | dd|j } | } || ddd|jf|d } t j| dd } | ddddf}t| d |ddd|jfd d } | jd|jdkr|rtrtt|j| }ntt|j| }|D]D}| dd|d|j|df}|||d } | dddddf} t j| dd }t|d|dd|fd}tj| |fdd } qt| } | dd|f} | |jdf| }|r$tj|jd d|d }||dddd f}|S)a Computes the log-likelihoods of the latent representations of the input. Args: inputs: input images, NxCxHxW[xD] vqvae_model: first stage model. transformer_model: autoregressive transformer model. ordering: ordering of the quantised latent representation. condition: conditioning for network input. resample_latent_likelihoods: if true, resamples the intermediate likelihood maps to have the same spatial dimension as the input images. resample_interpolation_mode: if use resample_latent_likelihoods, select interpolation 'nearest', 'bilinear', or 'trilinear; verbose: if true, prints the progression bar of the sampling process. rrNrErrrrrrrr.)r[rgryrrprrrrnrrrrrKrrgatherrrr r*rcr!rhrWrrr)r>r4rrrrrrrrrrrrrr(rjrpZprobs_reshapedrr?r?r@rjOsF    ,"   z&VQVAETransformerInferer.get_likelihoodr)NF)r4r5rrrrrr"rrrrRr9r)NrkNT)rrrr5rrrrrr"rrrrrrrrRr9r5)NFrF)r4r5rrrrrr"rrrrRrrZrrRr9r5) r<rBrCrDrVrArgryrmrjr?r?r?r@rs$   1?r)T __future__rrnrabcrrcollections.abcrrrrr functoolsr pydocr typingr rgtorch.nnrtorch.nn.functional functionalrmonai.apps.utilsr monai.datarmonai.data.meta_tensorrZmonai.data.thread_bufferrZmonai.inferers.mergerrrZmonai.inferers.splitterrZmonai.inferers.utilsrrmonai.networks.netsrrrrrrrZmonai.networks.schedulersrrmonai.transformsrr monai.utilsr!r"r#r$r%r&Zmonai.visualizer'r(r)r*r r<r__all__r-r.r/r0r3r1r2rrzrrModulerr?r?r?r@s`          $   &B#Y/jLb[