move load scheduler to method

hunyuan_train_network.py

@@ -5,6 +5,7 @@
 import argparse
 
 import torch
+from diffusers import DDPMScheduler
 from library.device_utils import init_ipex, clean_memory_on_device
 
 init_ipex()
@@ -77,6 +78,16 @@ def load_tokenizer(self, args):
         tokenizer = hunyuan_utils.load_tokenizers()
         return tokenizer
 
+    def load_noise_scheduler(self):
+        return DDPMScheduler(
+            beta_start=0.00085, 
+            beta_end=0.03, 
+            beta_schedule="scaled_linear", 
+            num_train_timesteps=1000, 
+            clip_sample=False,
+            steps_offset=1
+        )
+
     def is_text_encoder_outputs_cached(self, args):
         return args.cache_text_encoder_outputs
 

train_network.py

@@ -105,6 +105,12 @@ def load_tokenizer(self, args):
         tokenizer = train_util.load_tokenizer(args)
         return tokenizer
 
+    def load_noise_scheduler(self):
+        noise_scheduler = DDPMScheduler(
+            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
+        )
+        return noise_scheduler
+
     def is_text_encoder_outputs_cached(self, args):
         return False
 
@@ -839,9 +845,7 @@ def load_model_hook(models, input_dir):
 
         global_step = 0
 
-        noise_scheduler = DDPMScheduler(
-            beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, clip_sample=False
-        )
+        noise_scheduler = self.load_noise_scheduler()
         prepare_scheduler_for_custom_training(noise_scheduler, accelerator.device)
         if args.zero_terminal_snr:
             custom_train_functions.fix_noise_scheduler_betas_for_zero_terminal_snr(noise_scheduler)