(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
)
| 56 | |
| 57 | @torch.no_grad() |
| 58 | def sample(self, |
| 59 | S, |
| 60 | batch_size, |
| 61 | shape, |
| 62 | conditioning=None, |
| 63 | callback=None, |
| 64 | normals_sequence=None, |
| 65 | img_callback=None, |
| 66 | quantize_x0=False, |
| 67 | eta=0., |
| 68 | mask=None, |
| 69 | x0=None, |
| 70 | temperature=1., |
| 71 | noise_dropout=0., |
| 72 | score_corrector=None, |
| 73 | corrector_kwargs=None, |
| 74 | verbose=True, |
| 75 | x_T=None, |
| 76 | log_every_t=100, |
| 77 | unconditional_guidance_scale=1., |
| 78 | unconditional_conditioning=None, |
| 79 | # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... |
| 80 | **kwargs |
| 81 | ): |
| 82 | if conditioning is not None: |
| 83 | if isinstance(conditioning, dict): |
| 84 | cbs = conditioning[list(conditioning.keys())[0]].shape[0] |
| 85 | if cbs != batch_size: |
| 86 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}") |
| 87 | else: |
| 88 | if conditioning.shape[0] != batch_size: |
| 89 | print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") |
| 90 | |
| 91 | self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose) |
| 92 | # sampling |
| 93 | C, H, W = shape |
| 94 | size = (batch_size, C, H, W) |
| 95 | print(f'Data shape for PLMS sampling is {size}') |
| 96 | |
| 97 | samples, intermediates = self.plms_sampling(conditioning, size, |
| 98 | callback=callback, |
| 99 | img_callback=img_callback, |
| 100 | quantize_denoised=quantize_x0, |
| 101 | mask=mask, x0=x0, |
| 102 | ddim_use_original_steps=False, |
| 103 | noise_dropout=noise_dropout, |
| 104 | temperature=temperature, |
| 105 | score_corrector=score_corrector, |
| 106 | corrector_kwargs=corrector_kwargs, |
| 107 | x_T=x_T, |
| 108 | log_every_t=log_every_t, |
| 109 | unconditional_guidance_scale=unconditional_guidance_scale, |
| 110 | unconditional_conditioning=unconditional_conditioning, |
| 111 | **kwargs |
| 112 | ) |
| 113 | return samples, intermediates |
| 114 | |
| 115 | @torch.no_grad() |
no test coverage detected