class_docstr: Solver for MusicGen training task.
  
  Used in: https://arxiv.org/abs/2306.05284
docstr: Audio generation stage.
docstr_tok: ["Audio","generation","stage","."]
code: def generate_audio(self) -> dict:
      """Audio generation stage."""
      generate_stage_name = f"{self.current_stage}"
      sample_manager = SampleManager(self.xp)
      self.logger.info(
          f"Generating samples in {sample_manager.base_folder}"
      )
      loader = self.dataloaders["generate"]
      updates = len(loader)
      lp = self.log_progress(
          generate_stage_name,
          loader,
          total=updates,
          updates=self.log_updates,
      )
  
      dataset = get_dataset_from_loader(loader)
      dataset_duration = dataset.segment_duration
      assert dataset_duration is not None
      assert isinstance(dataset, AudioDataset)
      target_duration = self.cfg.generate.lm.gen_duration
      prompt_duration = self.cfg.generate.lm.prompt_duration
      if target_duration is None:
          target_duration = dataset_duration
      if prompt_duration is None:
          prompt_duration = dataset_duration / 4
      assert prompt_duration < dataset_duration, (
          f"Specified prompt duration ({prompt_duration}s) is longer",
          f" than reference audio duration ({dataset_duration}s)",
      )
  
      def get_hydrated_conditions(meta: tp.List[SegmentWithAttributes]):
          hydrated_conditions = []
          for sample in [x.to_condition_attributes() for x in meta]:
              cond_dict = {}
              for cond_type in sample.__annotations__.keys():
                  for cond_key, cond_val in getattr(
                      sample, cond_type
                  ).items():
                      if (
                          cond_key
                          not in self.model.condition_provider.conditioners.keys()
                      ):
                          continue
                      if is_jsonable(cond_val):
                          cond_dict[cond_key] = cond_val
                      elif isinstance(cond_val, WavCondition):
                          cond_dict[cond_key] = cond_val.path
                      elif isinstance(cond_val, JointEmbedCondition):
                          cond_dict[
                              cond_key
                          ] = (
                              cond_val.text
                          )  # only support text at inference for now
                      else:
                          # if we reached this point, it is not clear how to log the condition
                          # so we just log the type.
                          cond_dict[cond_key] = str(type(cond_val))
                          continue
              hydrated_conditions.append(cond_dict)
          return hydrated_conditions
  
      metrics: dict = {}
      average = flashy.averager()
      for batch in lp:
          audio, meta = batch
          # metadata for sample manager
          hydrated_conditions = get_hydrated_conditions(meta)
          sample_generation_params = {
              **{
                  f"classifier_free_guidance_{k}": v
                  for k, v in self.cfg.classifier_free_guidance.items()
              },
              **self.generation_params,
          }
          if self.cfg.generate.lm.unprompted_samples:
              if self.cfg.generate.lm.gen_gt_samples:
                  # get the ground truth instead of generation
                  self.logger.warn(
                      "Use ground truth instead of audio generation as generate.lm.gen_gt_samples=true"
                  )
                  gen_unprompted_audio = audio
                  rtf = 1.0
              else:
                  gen_unprompted_outputs = self.run_generate_step(
                      batch,
                      gen_duration=target_duration,
                      prompt_duration=None,
                      **self.generation_params,
                  )
                  gen_unprompted_audio = gen_unprompted_outputs[
                      "gen_audio"
                  ].cpu()
                  rtf = gen_unprompted_outputs["rtf"]
              sample_manager.add_samples(
                  gen_unprompted_audio,
                  self.epoch,
                  hydrated_conditions,
                  ground_truth_wavs=audio,
                  generation_args=sample_generation_params,
              )
  
          if self.cfg.generate.lm.prompted_samples:
              gen_outputs = self.run_generate_step(
                  batch,
                  gen_duration=target_duration,
                  prompt_duration=prompt_duration,
                  **self.generation_params,
              )
              gen_audio = gen_outputs["gen_audio"].cpu()
              prompt_audio = gen_outputs["prompt_audio"].cpu()
              sample_manager.add_samples(
                  gen_audio,
                  self.epoch,
                  hydrated_conditions,
                  prompt_wavs=prompt_audio,
                  ground_truth_wavs=audio,
                  generation_args=sample_generation_params,
              )
  
          metrics["rtf"] = rtf
          metrics = average(metrics)
  
      flashy.distrib.barrier()
      return metrics
  
imports: from pathlib import Path
  import time
  import typing
  import warnings
  import flashy
  import math
  import omegaconf
  import torch
  from torch.nn import functional
  from .None import base, builders
  from .compression import CompressionSolver
  from ..None import metrics
  from ..None import models
  from ..data.audio_dataset import AudioDataset
  from ..data.music_dataset import MusicDataset, MusicInfo, AudioInfo
  from ..data.audio_utils import normalize_audio
  from ..modules.conditioners import JointEmbedCondition, SegmentWithAttributes, WavCondition
  from ..utils.cache import CachedBatchWriter, CachedBatchLoader
  from ..utils.samples.manager import SampleManager
  from ..utils.utils import get_dataset_from_loader, is_jsonable, warn_once, model_hash
label_desc: Use image node_id 14 for calling the MusicGenSolver obj's underlying member method code with example usage: obj.generate_audio() and returns: metrics, hydrated_conditions