class_docstr: Performs Textual Inversion fine-tuning on a Determined cluster.
class_docstr_tok: ["Performs","Textual","Inversion","fine-tuning","on","a","Determined","cluster","."]
docstr: Train on a single batch and update internal metrics.
docstr_tok: ["Train","on","a","single","batch","and","update","internal","metrics","."]
code: def _train_one_batch(self, batch: TorchData) -> None:
      """Train on a single batch and update internal metrics."""
      self.text_encoder.train()
      # Convert sample images to latent space.
      prompts, img_tensors = batch
      with torch.no_grad():
          latent_dist = self.vae.encode(img_tensors).latent_dist
          latents = latent_dist.sample()
  
          # In 2112.10752, it was found that the latent space variance plays a large role in image
          # quality.  The following scale factor helps to maintain unit latent variance.  See
          # https://github.com/huggingface/diffusers/issues/437 for more details.
          scale_factor = 0.18215
          latents = latents * scale_factor
  
          # Sample noise that we'll add to the latents.
          noise = torch.randn(
              latents.shape, device=self.accelerator.device
          )
          # Sample a random timestep for each image in the batch.
          rand_timesteps = torch.randint(
              0,
              self.num_train_timesteps,
              (self.train_batch_size,),
              device=self.accelerator.device,
          ).long()
          # Add noise to the latents according to the noise magnitude at each timestep. This is
          # the forward diffusion process.
          noisy_latents = self.train_scheduler.add_noise(
              latents, noise, rand_timesteps
          )
  
      # Process the text for each batch.
      dummy_prompts = [
          self._replace_concepts_with_dummies(text) for text in prompts
      ]
      dummy_prompts_noise_pred = self._get_noise_pred(
          text=dummy_prompts,
          noisy_latents=noisy_latents,
          timesteps=rand_timesteps,
      )
  
      noise_pred_loss = F.mse_loss(dummy_prompts_noise_pred, noise)
      self.metrics_history["noise_pred_loss"].append(
          noise_pred_loss.item()
      )
      self.accelerator.backward(noise_pred_loss)
  
      # Optional regularization penalty for large learned embedding vectors.
      if self.norm_reg_weight:
          new_embedding = (
              self._get_token_embedding_layer().new_embedding.weight
          )
          new_embedding_norm = torch.linalg.norm(new_embedding, dim=1)
          new_embedding_norm_diff = (
              self.init_embedding_norm_mean - new_embedding_norm
          )
          norm_reg_loss = (
              self.norm_reg_weight
              * new_embedding_norm_diff.pow(2).sum()
          )
          self.metrics_history["norm_reg_loss"].append(
              norm_reg_loss.item()
          )
          self.accelerator.backward(norm_reg_loss)
  
      # Optional regularization penalty for changing the encoder hidden states too much.
      # Inspired by similar ideas in the original Textual Inversion repo and discussions on
      # the Stable Diffusion #community-research channel.
      if self.hidden_reg_weight:
          dummy_hidden_states = self._get_encoder_hidden_states(
              text=dummy_prompts
          )
          with torch.no_grad():
              intializer_prompts = [
                  self._replace_concepts_with_initializers(text)
                  for text in prompts
              ]
              initializer_hidden_states = (
                  self._get_encoder_hidden_states(
                      text=intializer_prompts
                  )
              )
          hidden_reg_loss = self.hidden_reg_weight * F.mse_loss(
              dummy_hidden_states, initializer_hidden_states
          )
          self.metrics_history["hidden_reg_loss"].append(
              hidden_reg_loss.item()
          )
          self.accelerator.backward(hidden_reg_loss)
  
      self.optimizer.step()
      self.optimizer.zero_grad()
  
imports: import json
  import os
  import pathlib
  from collections import defaultdict
  from typing import List, Literal, Optional, Sequence, Union
  import accelerate
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from detsd import data, defaults, layers, utils
  from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
  from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
  from torch.utils.data import DataLoader
  from torch.utils.tensorboard import SummaryWriter
  from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
  import determined
  from determined.pytorch import TorchData
label_desc: Use image node_id 3 for calling the DetSDTextualInversionTrainer obj's underlying member method code with example usage: obj._train_one_batch(batch) without return types