Args: x: source minibatch (bs, *dim) sample_kwargs: dict, additional sampling arguments for the solver progress: bool, whether to show a progress bar clip_denoised: bool, whether to clip the denoised images to [-1, 1] u
(self, x: torch.Tensor, sample_kwargs=None, reverse=False, return_intermediates=False, **kwargs)
| 114 | return loss |
| 115 | |
| 116 | def generate(self, x: torch.Tensor, sample_kwargs=None, reverse=False, return_intermediates=False, **kwargs): |
| 117 | """ |
| 118 | Args: |
| 119 | x: source minibatch (bs, *dim) |
| 120 | sample_kwargs: dict, additional sampling arguments for the solver |
| 121 | progress: bool, whether to show a progress bar |
| 122 | clip_denoised: bool, whether to clip the denoised images to [-1, 1] |
| 123 | use_ddpm: bool, whether to use DDPM sampling instead of DDIM |
| 124 | intermediate_key: str, key to use for intermediate outputs |
| 125 | (DDIM: 'x_inter', 'pred_x0' | DDPM: 'sample' or 'pred_xstart') |
| 126 | intermediate_freq: int, frequency of intermediate outputs |
| 127 | __ DDIM only __: |
| 128 | num_steps: int, number of DDIM steps to take |
| 129 | eta: float, noise level for DDIM |
| 130 | temperature: float, temperature for DDIM |
| 131 | noise_dropout: float, dropout rate for DDIM |
| 132 | cfg_scale: float, scale factor for Classifier-free guidance |
| 133 | uc_cond: torch.Tensor, unconditional conditioning information |
| 134 | cond_key: str, key to use for conditioning information |
| 135 | reverse: bool, whether to reverse the direction of the flow. If True, |
| 136 | we map from x1 -> x0, otherwise we map from x0 -> x1. |
| 137 | return_intermediates: if true, return the intermediate samples |
| 138 | kwargs: additional arguments for the network (e.g. conditioning information). |
| 139 | """ |
| 140 | if reverse: |
| 141 | raise NotImplementedError("[DiffusionFlow] Reverse sampling not yet supported") |
| 142 | |
| 143 | sample_kwargs = sample_kwargs or {} |
| 144 | |
| 145 | # DDPM sampling |
| 146 | if sample_kwargs.get("use_ddpm", False): |
| 147 | # include CFG |
| 148 | forward_fn = partial( |
| 149 | forward_with_cfg, |
| 150 | model = self.net, |
| 151 | cfg_scale = sample_kwargs.get("cfg_scale", 1.), |
| 152 | uc_cond = sample_kwargs.get("uc_cond", None), |
| 153 | cond_key = sample_kwargs.get("cond_key", "y"), |
| 154 | ) |
| 155 | out, intermediates = self.diffusion.p_sample_loop( |
| 156 | # model = self.net, # without CFG |
| 157 | model = forward_fn, |
| 158 | noise = x, |
| 159 | model_kwargs = kwargs, |
| 160 | progress = sample_kwargs.get("progress", False), |
| 161 | clip_denoised = sample_kwargs.get("clip_denoised", False), |
| 162 | return_intermediates = True, |
| 163 | intermediate_freq = sample_kwargs.get("intermediate_freq", (100 if return_intermediates else 1000)), |
| 164 | pbar_desc = sample_kwargs.get("pbar_desc", "DDPM Sampling"), |
| 165 | intermediate_key = sample_kwargs.get("intermediate_key", "sample"), |
| 166 | ) |
| 167 | |
| 168 | # DDIM sampling |
| 169 | else: |
| 170 | out, intermediates = self.ddim_sampler.sample( |
| 171 | model = self.net, |
| 172 | noise = x, |
| 173 | model_kwargs = kwargs, |
nothing calls this directly
no test coverage detected