Skip to content
This repository was archived by the owner on Feb 7, 2025. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion generative/networks/nets/autoencoderkl.py
Original file line number Diff line number Diff line change
Expand Up @@ -612,7 +612,9 @@ def __init__(
num_res_blocks = ensure_tuple_rep(num_res_blocks, len(num_channels))

if len(num_res_blocks) != len(num_channels):
raise ValueError("`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`.")
raise ValueError(
"`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`."
)

self.encoder = Encoder(
spatial_dims=spatial_dims,
Expand Down
4 changes: 3 additions & 1 deletion generative/networks/nets/diffusion_model_unet.py
Original file line number Diff line number Diff line change
Expand Up @@ -1655,7 +1655,9 @@ def __init__(
num_res_blocks = ensure_tuple_rep(num_res_blocks, len(num_channels))

if len(num_res_blocks) != len(num_channels):
raise ValueError("`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`.")
raise ValueError(
"`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`."
)

self.in_channels = in_channels
self.block_out_channels = num_channels
Expand Down
1 change: 0 additions & 1 deletion generative/networks/schedulers/ddim.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,6 @@ def __init__(
# standard deviation of the initial noise distribution
self.init_noise_sigma = 1.0


self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].astype(np.int64))

self.clip_sample = clip_sample
Expand Down
2 changes: 1 addition & 1 deletion generative/networks/schedulers/pndm.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,11 +117,11 @@ def __init__(
self.cur_sample = None
self.ets = []


self._timesteps = np.arange(0, num_train_timesteps)[::-1].copy()

# default the number of inference timesteps to the number of train steps
self.set_timesteps(num_train_timesteps)

def set_timesteps(self, num_inference_steps: int, device: str | torch.device | None = None) -> None:
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
Expand Down
1 change: 0 additions & 1 deletion tests/test_scheduler_pndm.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ def test_add_noise_2d_shape(self, input_param, input_shape, expected_shape):
noisy = scheduler.add_noise(original_samples=original_sample, noise=noise, timesteps=timesteps)
self.assertEqual(noisy.shape, expected_shape)


@parameterized.expand(TEST_CASES)
def test_step_shape(self, input_param, input_shape, expected_shape):
scheduler = PNDMScheduler(**input_param)
Expand Down