ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl → 0.2.0.dev20240531__py3-none-any.whl

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -0,0 +1,222 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import argparse
+import os
+from pathlib import Path
+from typing import Dict, Optional
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+import ai_edge_torch.generative.examples.stable_diffusion.samplers as samplers
+from ai_edge_torch.generative.examples.stable_diffusion.tokenizer import Tokenizer  # NOQA
+import ai_edge_torch.generative.examples.stable_diffusion.util as util
+from ai_edge_torch.model import TfLiteModel
+
+arg_parser = argparse.ArgumentParser()
+arg_parser.add_argument(
+    '--tokenizer_vocab_dir',
+    type=str,
+    help='Directory to the tokenizer vocabulary files, which include `merges.txt` and `vocab.json`',
+    required=True,
+)
+arg_parser.add_argument(
+    '--clip_ckpt', type=str, help='Path to CLIP TFLite tflite file', required=True
+)
+arg_parser.add_argument(
+    '--diffusion_ckpt', type=str, help='Path to diffusion tflite file', required=True
+)
+arg_parser.add_argument(
+    '--decoder_ckpt', type=str, help='Path to decoder tflite file', required=True
+)
+arg_parser.add_argument(
+    '--output_path',
+    type=str,
+    help='Path to the output generated image file.',
+    required=True,
+)
+arg_parser.add_argument(
+    '--prompt',
+    default='a photograph of an astronaut riding a horse',
+    type=str,
+    help='The prompt to guide the image generation.',
+)
+arg_parser.add_argument(
+    '--n_inference_steps', default=20, type=int, help='The number of denoising steps.'
+)
+arg_parser.add_argument(
+    '--sampler',
+    default='k_euler',
+    type=str,
+    choices=['k_euler', 'k_euler_ancestral', 'k_lms'],
+    help='A sampler to be used to denoise the encoded image latents. Can be one of `k_lms, `k_euler`, or `k_euler_ancestral`.',
+)
+
+
+class StableDiffusion:
+
+  def __init__(
+      self,
+      *,
+      tokenizer_vocab_dir: str,
+      clip_ckpt: str,
+      encoder_ckpt: Optional[str] = None,
+      diffusion_ckpt: str,
+      decoder_ckpt: str
+  ):
+    self.tokenizer = Tokenizer(tokenizer_vocab_dir)
+    self.clip = TfLiteModel.load(clip_ckpt)
+    self.decoder = TfLiteModel.load(decoder_ckpt)
+    self.diffusion = TfLiteModel.load(diffusion_ckpt)
+    if encoder_ckpt is not None:
+      self.encoder = TfLiteModel.load(encoder_ckpt)
+
+
+def run_tflite_pipeline(
+    model: StableDiffusion,
+    prompt: str,
+    output_path: str,
+    uncond_prompt: Optional[str] = None,
+    cfg_scale: float = 7.5,
+    height: int = 512,
+    width: int = 512,
+    sampler: str = 'k_euler',
+    n_inference_steps: int = 20,
+    seed: Optional[int] = None,
+    strength: float = 0.8,
+    input_image: Optional[Image.Image] = None,
+):
+  """Run stable diffusion pipeline with tflite model.
+  model:
+    StableDiffsuion model.
+  prompt:
+    The prompt to guide the image generation.
+  output_path:
+    The path to the generated output image.
+  uncond_prompt:
+    The prompt not to guide the image generation.
+  cfg_scale:
+    Guidance scale of classifier-free guidance. Higher guidance scale encourages to generate
+    images that are closely linked to the text `prompt`, usually at the expense of lower
+    image quality.
+  height:
+    The height in pixels of the generated image.
+  width:
+    The width in pixels of the generated image.
+  sampler:
+    A sampler to be used to denoise the encoded image latents. Can be one of `k_lms, `k_euler`,
+    or `k_euler_ancestral`.
+  n_inference_steps:
+    The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+    expense of slower inference. This parameter will be modulated by `strength`.
+  seed:
+    A seed to make generation deterministic.
+  strength:
+    Conceptually, indicates how much to transform the reference `input_image`. Must be between 0 and 1.
+    `input_image` will be used as a starting point, adding more noise to it the larger the `strength`.
+    The number of denoising steps depends on the amount of noise initially added. When `strength` is 1,
+    added noise will be maximum and the denoising process will run for the full number of iterations
+    specified in `n_inference_steps`. A value of 1, therefore, essentially ignores `input_image`.
+  input_image:
+    Image which is served as the starting point for the image generation.
+  """
+  if not 0 < strength < 1:
+    raise ValueError('strength must be between 0 and 1')
+  if height % 8 or width % 8:
+    raise ValueError('height and width must be a multiple of 8')
+  if seed is not None:
+    np.random.seed(seed)
+  if uncond_prompt is None:
+    uncond_prompt = ''
+
+  if sampler == 'k_lms':
+    sampler = samplers.KLMSSampler(n_inference_steps=n_inference_steps)
+  elif sampler == 'k_euler':
+    sampler = samplers.KEulerSampler(n_inference_steps=n_inference_steps)
+  elif sampler == 'k_euler_ancestral':
+    sampler = samplers.KEulerAncestralSampler(n_inference_steps=n_inference_steps)
+  else:
+    raise ValueError(
+        'Unknown sampler value %s. '
+        'Accepted values are {k_lms, k_euler, k_euler_ancestral}' % sampler
+    )
+
+  # Text embedding.
+  cond_tokens = model.tokenizer.encode(prompt)
+  cond_context = model.clip(np.array(cond_tokens), signature_name='encode')
+  uncond_tokens = model.tokenizer.encode(uncond_prompt)
+  uncond_context = model.clip(np.array(uncond_tokens), signature_name='encode')
+  context = np.concatenate([cond_context, uncond_context], axis=0)
+  noise_shape = (1, 4, height // 8, width // 8)
+
+  # Initialization starts from input_image if any, otherwise, starts from a random sampling.
+  if input_image:
+    if not hasattr(model, 'encoder'):
+      raise AttributeError(
+          'Stable Diffusion must be initilaized with encoder to accept input_image.'
+      )
+    input_image = input_image.resize((width, height))
+    input_image_np = np.array(input_image).astype(np.float32)
+    input_image_np = util.rescale(input_image, (0, 255), (-1, 1))
+    input_image_np = util.move_channel(input_image_np, to='first')
+    encoder_noise = np.random.normal(size=noise_shape).astype(np.float32)
+    latents = model.encoder(input_image_np, encoder_noise)
+    latents_noise = np.random.normal(size=noise_shape).astype(np.float32)
+    sampler.set_strength(strength=strength)
+    latents += latents_noise * sampler.initial_scale
+  else:
+    latents = np.random.normal(size=noise_shape).astype(np.float32)
+    latents *= sampler.initial_scale
+
+  # Diffusion process.
+  timesteps = tqdm(sampler.timesteps)
+  for i, timestep in enumerate(timesteps):
+    time_embedding = util.get_time_embedding(timestep)
+
+    input_latents = latents * sampler.get_input_scale()
+    input_latents = input_latents.repeat(2, axis=0)
+    output = model.diffusion(
+        input_latents, context, time_embedding, signature_name='diffusion'
+    )
+    output_cond, output_uncond = np.split(output, 2, axis=0)
+    output = cfg_scale * (output_cond - output_uncond) + output_uncond
+
+    latents = sampler.step(latents, output)
+
+  # Image decoding.
+  images = model.decoder(latents, signature_name='decode')
+  images = util.rescale(images, (-1, 1), (0, 255), clamp=True)
+  images = util.move_channel(images, to='last')
+  if not os.path.exists(output_path):
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+  Image.fromarray(images[0].astype(np.uint8)).save(output_path)
+
+
+if __name__ == '__main__':
+  args = arg_parser.parse_args()
+  run_tflite_pipeline(
+      StableDiffusion(
+          tokenizer_vocab_dir=args.tokenizer_vocab_dir,
+          clip_ckpt=args.clip_ckpt,
+          diffusion_ckpt=args.diffusion_ckpt,
+          decoder_ckpt=args.decoder_ckpt,
+      ),
+      prompt=args.prompt,
+      output_path=args.output_path,
+      sampler=args.sampler,
+      n_inference_steps=args.n_inference_steps,
+  )

ai_edge_torch/generative/examples/stable_diffusion/samplers/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from .k_euler import KEulerSampler
+from .k_euler_ancestral import KEulerAncestralSampler
+from .k_lms import KLMSSampler
+from .sampler import SamplerInterface

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler.py

@@ -0,0 +1,61 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import numpy as np
+
+from ai_edge_torch.generative.examples.stable_diffusion import util
+from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+
+
+class KEulerSampler(SamplerInterface):
+
+  def __init__(self, n_inference_steps=50, n_training_steps=1000):
+    timesteps = np.linspace(n_training_steps - 1, 0, n_inference_steps)
+
+    alphas_cumprod = util.get_alphas_cumprod(n_training_steps=n_training_steps)
+    sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+    log_sigmas = np.log(sigmas)
+    log_sigmas = np.interp(timesteps, range(n_training_steps), log_sigmas)
+    sigmas = np.exp(log_sigmas)
+    sigmas = np.append(sigmas, 0)
+
+    self.sigmas = sigmas
+    self.initial_scale = sigmas.max()
+    self.timesteps = timesteps
+    self.n_inference_steps = n_inference_steps
+    self.n_training_steps = n_training_steps
+    self.step_count = 0
+
+  def get_input_scale(self, step_count=None):
+    if step_count is None:
+      step_count = self.step_count
+    sigma = self.sigmas[step_count]
+    return 1 / (sigma**2 + 1) ** 0.5
+
+  def set_strength(self, strength=1):
+    start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
+    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = self.timesteps[start_step:]
+    self.initial_scale = self.sigmas[start_step]
+    self.step_count = start_step
+
+  def step(self, latents, output):
+    t = self.step_count
+    self.step_count += 1
+
+    sigma_from = self.sigmas[t]
+    sigma_to = self.sigmas[t + 1]
+    latents += output * (sigma_to - sigma_from)
+    return latents

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_euler_ancestral.py

@@ -0,0 +1,65 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import numpy as np
+
+from ai_edge_torch.generative.examples.stable_diffusion import util
+from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+
+
+class KEulerAncestralSampler(SamplerInterface):
+
+  def __init__(self, n_inference_steps=50, n_training_steps=1000):
+    timesteps = np.linspace(n_training_steps - 1, 0, n_inference_steps)
+
+    alphas_cumprod = util.get_alphas_cumprod(n_training_steps=n_training_steps)
+    sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+    log_sigmas = np.log(sigmas)
+    log_sigmas = np.interp(timesteps, range(n_training_steps), log_sigmas)
+    sigmas = np.exp(log_sigmas)
+    sigmas = np.append(sigmas, 0)
+
+    self.sigmas = sigmas
+    self.initial_scale = sigmas.max()
+    self.timesteps = timesteps
+    self.n_inference_steps = n_inference_steps
+    self.n_training_steps = n_training_steps
+    self.step_count = 0
+
+  def get_input_scale(self, step_count=None):
+    if step_count is None:
+      step_count = self.step_count
+    sigma = self.sigmas[step_count]
+    return 1 / (sigma**2 + 1) ** 0.5
+
+  def set_strength(self, strength=1):
+    start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
+    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = self.timesteps[start_step:]
+    self.initial_scale = self.sigmas[start_step]
+    self.step_count = start_step
+
+  def step(self, latents, output):
+    t = self.step_count
+    self.step_count += 1
+
+    sigma_from = self.sigmas[t]
+    sigma_to = self.sigmas[t + 1]
+    sigma_up = sigma_to * (1 - (sigma_to**2 / sigma_from**2)) ** 0.5
+    sigma_down = sigma_to**2 / sigma_from
+    latents += output * (sigma_down - sigma_from)
+    noise = np.random.normal(size=latents.shape)
+    latents += noise * sigma_up
+    return latents

ai_edge_torch/generative/examples/stable_diffusion/samplers/k_lms.py

@@ -0,0 +1,73 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import numpy as np
+
+from ai_edge_torch.generative.examples.stable_diffusion import util
+from ai_edge_torch.generative.examples.stable_diffusion.samplers.sampler import SamplerInterface  # NOQA
+
+
+class KLMSSampler(SamplerInterface):
+
+  def __init__(self, n_inference_steps=50, n_training_steps=1000, lms_order=4):
+    timesteps = np.linspace(n_training_steps - 1, 0, n_inference_steps)
+
+    alphas_cumprod = util.get_alphas_cumprod(n_training_steps=n_training_steps)
+    sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
+    log_sigmas = np.log(sigmas)
+    log_sigmas = np.interp(timesteps, range(n_training_steps), log_sigmas)
+    sigmas = np.exp(log_sigmas)
+    sigmas = np.append(sigmas, 0)
+
+    self.sigmas = sigmas
+    self.initial_scale = sigmas.max()
+    self.timesteps = timesteps
+    self.n_inference_steps = n_inference_steps
+    self.n_training_steps = n_training_steps
+    self.lms_order = lms_order
+    self.step_count = 0
+    self.outputs = []
+
+  def get_input_scale(self, step_count=None):
+    if step_count is None:
+      step_count = self.step_count
+    sigma = self.sigmas[step_count]
+    return 1 / (sigma**2 + 1) ** 0.5
+
+  def set_strength(self, strength=1):
+    start_step = self.n_inference_steps - int(self.n_inference_steps * strength)
+    self.timesteps = np.linspace(self.n_training_steps - 1, 0, self.n_inference_steps)
+    self.timesteps = self.timesteps[start_step:]
+    self.initial_scale = self.sigmas[start_step]
+    self.step_count = start_step
+
+  def step(self, latents, output):
+    t = self.step_count
+    self.step_count += 1
+
+    self.outputs = [output] + self.outputs[: self.lms_order - 1]
+    order = len(self.outputs)
+    for i, output in enumerate(self.outputs):
+      # Integrate polynomial by trapezoidal approx. method for 81 points.
+      x = np.linspace(self.sigmas[t], self.sigmas[t + 1], 81)
+      y = np.ones(81)
+      for j in range(order):
+        if i == j:
+          continue
+        y *= x - self.sigmas[t - j]
+        y /= self.sigmas[t - i] - self.sigmas[t - j]
+      lms_coeff = np.trapz(y=y, x=x)
+      latents += lms_coeff * output
+    return latents

ai_edge_torch/generative/examples/stable_diffusion/samplers/sampler.py

@@ -0,0 +1,38 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import abc
+
+import numpy as np
+
+
+class SamplerInterface(abc.ABC):
+
+  @abc.abstractmethod
+  def get_input_scale(self, step_count: int = 1) -> float:
+    """Get the input scale of the random samples from sampled distribution"""
+    return NotImplemented
+
+  @abc.abstractmethod
+  def set_strength(self, strength: float = 1) -> None:
+    """Set the strength of initial step.
+    Conceptually, indicates how much to transform the reference `input_images`.
+    """
+    return NotImplemented
+
+  @abc.abstractmethod
+  def step(self, latents: np.ndarray, output: np.ndarray) -> np.ndarray:
+    """Update latents from the diffusion output by a step"""
+    return NotImplemented

ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py

@@ -0,0 +1,108 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import functools
+import json
+import os
+from typing import List, Tuple
+import unicodedata
+
+import regex as re
+
+
+def create_bytes_table() -> dict:
+  table = {}
+  special_count = 0
+  for byte in range(256):
+    category = unicodedata.category(chr(byte))
+    if category[0] not in ['C', 'Z']:  # ith character is NOT control char or space
+      table[byte] = chr(byte)
+    else:  # ith character IS control char or space
+      table[byte] = chr(special_count + 256)
+      special_count += 1
+  return table
+
+
+def pairwise(seq):
+  a = iter(seq)
+  b = iter(seq)
+  next(b)
+  return zip(a, b)
+
+
+class Tokenizer:
+
+  def __init__(self, vocab_dir: str):
+    with open(os.path.join(vocab_dir, 'vocab.json'), encoding='utf-8') as f:
+      self.vocab = json.load(f)
+
+    with open(os.path.join(vocab_dir, 'merges.txt'), encoding='utf-8') as f:
+      lines = f.read().split('\n')
+      lines = lines[1:-1]
+      self.merges = {tuple(bigram.split()): i for i, bigram in enumerate(lines)}
+
+    self.bos_token = self.vocab['<|startoftext|>']
+    self.eos_token = self.vocab['<|endoftext|>']
+    self.pad_token = self.vocab['<|endoftext|>']
+    self.max_length = 77
+    self.bytes_table = create_bytes_table()
+    self.chunk_pattern = re.compile(
+        r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+        re.IGNORECASE,
+    )
+
+  def encode(self, text: str) -> List[int]:
+    text = unicodedata.normalize('NFC', text)
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    text = text.lower()
+
+    tokens = [self.bos_token]
+    for chunk in re.findall(self.chunk_pattern, text):
+      chunk = ''.join(self.bytes_table[byte] for byte in chunk.encode('utf-8'))
+      tokens.extend(self.vocab[word] for word in self.bpe(chunk))
+    tokens.append(self.eos_token)
+
+    tokens = tokens[: self.max_length]
+    token_length = len(tokens)
+    pad_length = self.max_length - token_length
+    tokens += [self.pad_token] * pad_length
+    return tokens
+
+  def encode_batch(self, texts: List[str]) -> List[List[int]]:
+    return [self.encode(text) for text in texts]
+
+  @functools.lru_cache(maxsize=10000)
+  def bpe(self, chunk: str) -> Tuple[str]:
+    words = list(chunk)
+    words[-1] += '</w>'
+
+    while len(words) > 1:
+      valid_pairs = [pair for pair in pairwise(words) if pair in self.merges]
+      if not valid_pairs:
+        break
+
+      bigram = min(valid_pairs, key=lambda pair: self.merges[pair])
+      first, second = bigram
+
+      new_words = []
+      for word in words:
+        if word == second and new_words and new_words[-1] == first:
+          new_words[-1] = first + second
+        else:
+          new_words.append(word)
+      words = new_words
+
+    return tuple(words)

ai_edge_torch/generative/examples/stable_diffusion/util.py

@@ -0,0 +1,71 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+
+import numpy as np
+import torch
+
+
+def get_time_embedding(timestep):
+  freqs = torch.pow(10000, -torch.arange(start=0, end=160, dtype=torch.float32) / 160)
+  x = torch.tensor([timestep], dtype=torch.float32)[:, None] * freqs[None]
+  return torch.cat([torch.cos(x), torch.sin(x)], dim=-1)
+
+
+def get_alphas_cumprod(beta_start=0.00085, beta_end=0.0120, n_training_steps=1000):
+  betas = (
+      np.linspace(beta_start**0.5, beta_end**0.5, n_training_steps, dtype=np.float32)
+      ** 2
+  )
+  alphas = 1.0 - betas
+  alphas_cumprod = np.cumprod(alphas, axis=0)
+  return alphas_cumprod
+
+
+def get_file_path(filename, url=None):
+  module_location = os.path.dirname(os.path.abspath(__file__))
+  parent_location = os.path.dirname(module_location)
+  file_location = os.path.join(parent_location, "data", filename)
+  return file_location
+
+
+def move_channel(image, to):
+  if to == "first":
+    if isinstance(image, torch.Tensor):
+      return image.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+    if isinstance(image, np.ndarray):
+      return image.transpose(0, 3, 1, 2)
+  elif to == "last":
+    if isinstance(image, torch.Tensor):
+      return image.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+    if isinstance(image, np.ndarray):
+      return image.transpose(0, 2, 3, 1)
+  else:
+    raise ValueError("to must be one of the following: first, last")
+
+
+def rescale(x, old_range, new_range, clamp=False):
+  old_min, old_max = old_range
+  new_min, new_max = new_range
+  x -= old_min
+  x *= (new_max - new_min) / (old_max - old_min)
+  x += new_min
+  if clamp:
+    if isinstance(x, torch.Tensor):
+      x = x.clamp(new_min, new_max)
+    elif isinstance(x, np.ndarray):
+      x = x.clip(new_min, new_max)
+  return x