1

Author: zRzRzRzRzRzRzR

examples/cogview4-control/train_control_cogview4.py

@@ -1025,18 +1025,6 @@ def load_model_hook(models, input_dir):
         disable=not accelerator.is_local_main_process,
     )
 
-    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
-        sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
-        schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
-        timesteps = timesteps.to(accelerator.device)
-        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
-
-        sigma = sigmas[step_indices].flatten()
-        while len(sigma.shape) < n_dim:
-            sigma = sigma.unsqueeze(-1)
-        return sigma
-
-    image_logs = None
     for epoch in range(first_epoch, args.num_train_epochs):
         cogview4_transformer.train()
         for step, batch in enumerate(train_dataloader):