diff --git a/MindFlow/mindflow/cell/diffusion.py b/MindFlow/mindflow/cell/diffusion.py index ee2b8b67d084322de54546f724afac4f81e4f725..7444a9104a8932701a87b34df62d7572b09c9da9 100644 --- a/MindFlow/mindflow/cell/diffusion.py +++ b/MindFlow/mindflow/cell/diffusion.py @@ -591,7 +591,7 @@ class DDIMScheduler(DiffusionScheduler): Examples: >>> from mindspore import ops, dtype as mstype - >>> from mindflow.cell import DDPMScheduler + >>> from mindflow.cell import DDIMScheduler >>> scheduler = DDIMScheduler(num_train_timesteps=1000, ... beta_start=0.0001, ... beta_end=0.02, @@ -812,7 +812,7 @@ class DiffusionPipeline: ... clip_sample=True, ... clip_sample_range=1.0, ... thresholding=False, - ... dynamic_thresholding_ratio=None, + ... dynamic_thresholding_ratio=0.995, ... rescale_betas_zero_snr=False, ... timestep_spacing="leading", ... compute_dtype=mstype.float32) @@ -918,7 +918,7 @@ class DDPMPipeline(DiffusionPipeline): ... clip_sample=True, ... clip_sample_range=1.0, ... thresholding=False, - ... dynamic_thresholding_ratio=None, + ... dynamic_thresholding_ratio=0.995, ... rescale_betas_zero_snr=False, ... timestep_spacing="leading", ... compute_dtype=mstype.float32) @@ -962,7 +962,7 @@ class DDIMPipeline(DiffusionPipeline): Examples: >>> from mindspore import ops, dtype as mstype - >>> from mindflow.cell import DDIMPipeline, DDPMScheduler, ConditionDiffusionTransformer + >>> from mindflow.cell import DDIMPipeline, DDIMScheduler, ConditionDiffusionTransformer >>> # init params >>> in_dim, out_dim, hidden_dim, cond_dim, layers, heads, seq_len, batch_size = 16, 16, 256, 4, 3, 4, 256, 8 >>> # init condition @@ -977,14 +977,16 @@ class DDIMPipeline(DiffusionPipeline): ... time_token_cond=True, ... compute_dtype=mstype.float32) >>> num_train_timesteps = 100 - >>> scheduler = DDPMScheduler(num_train_timesteps=num_train_timesteps, + >>> scheduler = DDIMScheduler(num_train_timesteps=num_train_timesteps, ... beta_start=0.0001, ... beta_end=0.02, ... beta_schedule="squaredcos_cap_v2", + ... prediction_type='epsilon', ... clip_sample=True, ... clip_sample_range=1.0, ... thresholding=False, - ... dynamic_thresholding_ratio=None, + ... sample_max_value=1., + ... dynamic_thresholding_ratio=0.995, ... rescale_betas_zero_snr=False, ... timestep_spacing="leading", ... compute_dtype=mstype.float32) @@ -1093,7 +1095,7 @@ class DiffusionTrainer: ... clip_sample=True, ... clip_sample_range=1.0, ... thresholding=False, - ... dynamic_thresholding_ratio=None, + ... dynamic_thresholding_ratio=0.995, ... rescale_betas_zero_snr=False, ... timestep_spacing="leading", ... compute_dtype=mstype.float32) diff --git a/MindFlow/mindflow/cell/diffusion_transformer.py b/MindFlow/mindflow/cell/diffusion_transformer.py index b8318146f8eef958f196a1d2b2de41c33cb384f9..98e94bd49713851705ded2918731d9e1bbf61330 100644 --- a/MindFlow/mindflow/cell/diffusion_transformer.py +++ b/MindFlow/mindflow/cell/diffusion_transformer.py @@ -218,14 +218,15 @@ class ConditionDiffusionTransformer(DiffusionTransformer): Examples: >>> from mindspore import ops - >>> from mindflow.cell import DiffusionTransformer + >>> from mindflow.cell import ConditionDiffusionTransformer >>> in_channels, out_channels, cond_channels, hidden_channels = 16, 16, 10, 256 >>> layers, heads, batch_size, seq_len = 3, 4, 8, 256 - >>> model = DiffusionTransformer(in_channels=in_channels, - ... out_channels=out_channels, - ... hidden_channels=hidden_channels, - ... layers=layers, - ... heads=heads) + >>> model = ConditionDiffusionTransformer(in_channels=in_channels, + ... out_channels=out_channels, + ... cond_channels=cond_channels, + ... hidden_channels=hidden_channels, + ... layers=layers, + ... heads=heads) >>> x = ops.rand((batch_size, seq_len, in_channels)) >>> cond = ops.rand((batch_size, cond_channels)) >>> timestep = ops.randint(0, 1000, (batch_size,))