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