(self,
S,
batch_size,
shape,
conditioning=None,
callback=None,
normals_sequence=None,
img_callback=None,
quantize_x0=False,
eta=0.,
mask=None,
x0=None,
temperature=1.,
noise_dropout=0.,
score_corrector=None,
corrector_kwargs=None,
verbose=True,
x_T=None,
log_every_t=100,
unconditional_guidance_scale=1.,
unconditional_conditioning=None,
# this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
**kwargs
)
| 119 | |
| 120 | @torch.no_grad() |
| 121 | def sample(self, |
| 122 | S, |
| 123 | batch_size, |
| 124 | shape, |
| 125 | conditioning=None, |
| 126 | callback=None, |
| 127 | normals_sequence=None, |
| 128 | img_callback=None, |
| 129 | quantize_x0=False, |
| 130 | eta=0., |
| 131 | mask=None, |
| 132 | x0=None, |
| 133 | temperature=1., |
| 134 | noise_dropout=0., |
| 135 | score_corrector=None, |
| 136 | corrector_kwargs=None, |
| 137 | verbose=True, |
| 138 | x_T=None, |
| 139 | log_every_t=100, |
| 140 | unconditional_guidance_scale=1., |
| 141 | unconditional_conditioning=None, |
| 142 | # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... |
| 143 | **kwargs |
| 144 | ): |
| 145 | if conditioning is not None: |
| 146 | if isinstance(conditioning, dict): |
| 147 | cbs = conditioning[list(conditioning.keys())[0]].shape[0] |
| 148 | if cbs != batch_size: |
| 149 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") |
| 150 | else: |
| 151 | if conditioning.shape[0] != batch_size: |
| 152 | print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") |
| 153 | |
| 154 | self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose) |
| 155 | # sampling |
| 156 | C, H, W = shape |
| 157 | size = (batch_size, C, H, W) |
| 158 | print(f'Data shape for DDIM sampling is {size}, eta {eta}') |
| 159 | |
| 160 | samples, intermediates = self.ddim_sampling(conditioning, size, |
| 161 | callback=callback, |
| 162 | img_callback=img_callback, |
| 163 | quantize_denoised=quantize_x0, |
| 164 | mask=mask, x0=x0, |
| 165 | ddim_use_original_steps=False, |
| 166 | noise_dropout=noise_dropout, |
| 167 | temperature=temperature, |
| 168 | score_corrector=score_corrector, |
| 169 | corrector_kwargs=corrector_kwargs, |
| 170 | x_T=x_T, |
| 171 | log_every_t=log_every_t, |
| 172 | unconditional_guidance_scale=unconditional_guidance_scale, |
| 173 | unconditional_conditioning=unconditional_conditioning, |
| 174 | ) |
| 175 | return samples, intermediates |
| 176 | |
| 177 | @torch.no_grad() |
| 178 | def ddim_sampling(self, cond, shape, |
no test coverage detected