|
import torch |
|
import torch.nn.functional as F |
|
|
|
from diffusers import VQDiffusionScheduler |
|
|
|
from .test_schedulers import SchedulerCommonTest |
|
|
|
|
|
class VQDiffusionSchedulerTest(SchedulerCommonTest): |
|
scheduler_classes = (VQDiffusionScheduler,) |
|
|
|
def get_scheduler_config(self, **kwargs): |
|
config = { |
|
"num_vec_classes": 4097, |
|
"num_train_timesteps": 100, |
|
} |
|
|
|
config.update(**kwargs) |
|
return config |
|
|
|
def dummy_sample(self, num_vec_classes): |
|
batch_size = 4 |
|
height = 8 |
|
width = 8 |
|
|
|
sample = torch.randint(0, num_vec_classes, (batch_size, height * width)) |
|
|
|
return sample |
|
|
|
@property |
|
def dummy_sample_deter(self): |
|
assert False |
|
|
|
def dummy_model(self, num_vec_classes): |
|
def model(sample, t, *args): |
|
batch_size, num_latent_pixels = sample.shape |
|
logits = torch.rand((batch_size, num_vec_classes - 1, num_latent_pixels)) |
|
return_value = F.log_softmax(logits.double(), dim=1).float() |
|
return return_value |
|
|
|
return model |
|
|
|
def test_timesteps(self): |
|
for timesteps in [2, 5, 100, 1000]: |
|
self.check_over_configs(num_train_timesteps=timesteps) |
|
|
|
def test_num_vec_classes(self): |
|
for num_vec_classes in [5, 100, 1000, 4000]: |
|
self.check_over_configs(num_vec_classes=num_vec_classes) |
|
|
|
def test_time_indices(self): |
|
for t in [0, 50, 99]: |
|
self.check_over_forward(time_step=t) |
|
|
|
def test_add_noise_device(self): |
|
pass |
|
|
