From 7575cda0ab90db9f7d96579dc2dca3f4244c5754 Mon Sep 17 00:00:00 2001
From: brian <guoboqiang@huawei.com>
Date: Wed, 23 Apr 2025 17:04:55 +0800
Subject: [PATCH] [mod]: doc string

---
 MindFlow/mindflow/cell/diffusion.py             | 16 +++++++++-------
 MindFlow/mindflow/cell/diffusion_transformer.py | 13 +++++++------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/MindFlow/mindflow/cell/diffusion.py b/MindFlow/mindflow/cell/diffusion.py
index ee2b8b67d..7444a9104 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 b8318146f..98e94bd49 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,))
-- 
Gitee