class_docstr: Solver for compression task.
  
  The diffusion task allows for MultiBand diffusion model training.
  
  Args:
      cfg (DictConfig): Configuration.
class_docstr_tok: ["Solver","for","compression","task",".","The","diffusion","task","allows","for","MultiBand","diffusion","model","training",".","Args",":","cfg","(","DictConfig",")",":","Configuration","."]
docstr: Perform one training or valid step on a given batch.
docstr_tok: ["Perform","one","training","or","valid","step","on","a","given","batch","."]
code: def run_step(self, idx: int, batch: torch.Tensor, metrics: dict):
      """Perform one training or valid step on a given batch."""
      x = batch.to(self.device)
      loss_fun = (
          F.mse_loss if self.cfg.loss.kind == "mse" else F.l1_loss
      )
  
      condition = self.get_condition(x)  # [bs, 128, T/hop, n_emb]
      sample = self.data_processor.process_data(x)
  
      input_, target, step = self.schedule.get_training_item(
          sample, tensor_step=self.cfg.schedule.variable_step_batch
      )
      out = self.model(input_, step, condition=condition).sample
  
      base_loss = loss_fun(out, target, reduction="none").mean(
          dim=(1, 2)
      )
      reference_loss = loss_fun(input_, target, reduction="none").mean(
          dim=(1, 2)
      )
      loss = base_loss / reference_loss**self.cfg.loss.norm_power
  
      if self.is_training:
          loss.mean().backward()
          flashy.distrib.sync_model(self.model)
          self.optimizer.step()
          self.optimizer.zero_grad()
      metrics = {
          "loss": loss.mean(),
          "normed_loss": (base_loss / reference_loss).mean(),
      }
      metrics.update(
          self.per_stage(
              {"loss": loss, "normed_loss": base_loss / reference_loss},
              step,
          )
      )
      metrics.update({"std_in": input_.std(), "std_out": out.std()})
      return metrics
  
imports: import typing
  import flashy
  import julius
  import omegaconf
  import torch
  import torch.nn.functional
  from .None import builders
  from .None import base
  from ..None import models
  from ..modules.diffusion_schedule import NoiseSchedule
  from ..metrics import RelativeVolumeMel
  from ..models.builders import get_processor
  from ..utils.samples.manager import SampleManager
  from ..solvers.compression import CompressionSolver
label_desc: Use image node_id 7 for calling the DiffusionSolver obj's underlying member method code with example usage: obj.run_step(idx, batch, metrics) and returns: metrics