code: def compute_snr(noise_scheduler, timesteps):
      """
      Computes SNR as per
      https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
      """
      alphas_cumprod = noise_scheduler.alphas_cumprod
      sqrt_alphas_cumprod = alphas_cumprod**0.5
      sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
  
      # Expand the tensors.
      # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
      sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(
          device=timesteps.device
      )[timesteps].float()
      while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
          sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
      alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
  
      sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(
          device=timesteps.device
      )[timesteps].float()
      while len(sqrt_one_minus_alphas_cumprod.shape) < len(
          timesteps.shape
      ):
          sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[
              ..., None
          ]
      sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
  
      # Compute SNR.
      snr = (alpha / sigma) ** 2
      return snr
  
imports: import contextlib
  import copy
  import random
  from typing import Any, Dict, Iterable, List, Optional, Union
  import numpy
  import torch
  from torchvision import transforms
  from .models import UNet2DConditionModel
  from .utils import convert_state_dict_to_diffusers, convert_state_dict_to_peft, deprecate, is_peft_available, is_transformers_available
label_desc: Use image node_id 2 for calling a global function with example usage: compute_snr(noise_scheduler, timesteps) and returns: snr