prototype_general_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from flaxdiff.schedulers import EDMNoiseScheduler, KarrasVENoiseScheduler\n",
    "from flaxdiff.predictors import KarrasPredictionTransform\n",
    "from flaxdiff.models.simple_unet import Unet\n",
    "from flaxdiff.trainer.general_diffusion_trainer import GeneralDiffusionTrainer, ConditionalInputConfig\n",
    "from flaxdiff.data.dataloaders import get_dataset_grain\n",
    "from flaxdiff.utils import defaultTextEncodeModel, get_latest_checkpoint\n",
    "from flaxdiff.models.autoencoder.diffusers import StableDiffusionVAE\n",
    "from flaxdiff.samplers.euler import EulerAncestralSampler\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import optax\n",
    "from datetime import datetime\n",
    "import argparse\n",
    "import os\n",
    "\n",
    "BATCH_SIZE = 16\n",
    "IMAGE_SIZE = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-04-19 19:46:43.910069: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "E0000 00:00:1745092003.933578 2058403 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "E0000 00:00:1745092003.940782 2058403 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "W0000 00:00:1745092003.957749 2058403 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1745092003.957772 2058403 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1745092003.957774 2058403 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1745092003.957776 2058403 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "Some weights of the model checkpoint at openai/clip-vit-large-patch14 were not used when initializing FlaxCLIPTextModel: {('vision_model', 'encoder', 'layers', '13', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '1', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '0', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '6', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '16', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '0', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'layer_norm2', 'bias'), ('visual_projection', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '16', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '8', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '12', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '7', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '22', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '21', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '8', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '13', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '21', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '15', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '8', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '1', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '0', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '8', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '12', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '14', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '14', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '23', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '14', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '0', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '3', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '13', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '20', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '12', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '22', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '22', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '7', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '22', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '21', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '15', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '3', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '20', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'embeddings', 'position_embedding', 'embedding'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '6', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'embeddings', 'class_embedding'), ('vision_model', 'encoder', 'layers', '5', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'post_layernorm', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'q_proj', 'kernel'), ('logit_scale',), ('vision_model', 'encoder', 'layers', '8', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '4', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '23', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'post_layernorm', 'scale'), ('vision_model', 'encoder', 'layers', '1', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '2', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '13', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '14', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '4', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '4', 'layer_norm1', 'scale'), ('text_projection', 'kernel'), ('vision_model', 'encoder', 'layers', '5', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '10', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '7', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '7', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '11', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '23', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'embeddings', 'patch_embedding', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '19', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '13', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '17', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '7', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '13', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '3', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '6', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'pre_layrnorm', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '16', 'self_attn', 'k_proj', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '4', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '3', 'mlp', 'fc2', 'bias'), ('vision_model', 'encoder', 'layers', '22', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '11', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '23', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '3', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '0', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '2', 'self_attn', 'q_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '19', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '9', 'self_attn', 'out_proj', 'bias'), ('vision_model', 'encoder', 'layers', '12', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'mlp', 'fc1', 'kernel'), ('vision_model', 'encoder', 'layers', '15', 'self_attn', 'q_proj', 'bias'), ('vision_model', 'encoder', 'layers', '8', 'self_attn', 'k_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '6', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '20', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '18', 'layer_norm1', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'pre_layrnorm', 'scale'), ('vision_model', 'encoder', 'layers', '3', 'layer_norm2', 'scale'), ('vision_model', 'encoder', 'layers', '0', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'v_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '7', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'mlp', 'fc2', 'kernel'), ('vision_model', 'encoder', 'layers', '16', 'layer_norm2', 'bias'), ('vision_model', 'encoder', 'layers', '1', 'mlp', 'fc1', 'bias'), ('vision_model', 'encoder', 'layers', '21', 'self_attn', 'v_proj', 'bias'), ('vision_model', 'encoder', 'layers', '18', 'self_attn', 'out_proj', 'kernel'), ('vision_model', 'encoder', 'layers', '17', 'layer_norm1', 'scale'), ('vision_model', 'encoder', 'layers', '22', 'layer_norm1', 'bias')}\n",
      "- This IS expected if you are initializing FlaxCLIPTextModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing FlaxCLIPTextModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/torch_xla/__init__.py:251: UserWarning: `tensorflow` can conflict with `torch-xla`. Prefer `tensorflow-cpu` when using PyTorch/XLA. To silence this warning, `pip uninstall -y tensorflow && pip install tensorflow-cpu`. If you are in a notebook environment such as Colab or Kaggle, restart your notebook runtime afterwards.\n",
      "  warnings.warn(\n",
      "WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scaling factor: 0.18215\n",
      "Calculating downscale factor...\n"
     ]
    }
   ],
   "source": [
    "# Load dataset\n",
    "data = get_dataset_grain(\n",
    "    \"oxford_flowers102\", batch_size=BATCH_SIZE, image_scale=IMAGE_SIZE)\n",
    "datalen = data['train_len']\n",
    "batches = datalen // BATCH_SIZE\n",
    "\n",
    "text_encoder = defaultTextEncodeModel()\n",
    "autoencoder = StableDiffusionVAE(**{\"modelname\": \"pcuenq/sd-vae-ft-mse-flax\"})\n",
    "\n",
    "# Construct a validation set by the prompts\n",
    "val_prompts = ['water tulip', ' a water lily', ' a water lily', ' a photo of a rose', ' a photo of a rose', ' a water lily', ' a water lily', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a marigold', ' a water lily', ' a photo of a sunflower', ' a photo of a lotus', ' columbine', ' columbine', ' an orchid', ' an orchid', ' an orchid', ' a water lily', ' a water lily', ' a water lily', ' columbine', ' columbine', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a sunflower', ' a photo of a lotus', ' a photo of a lotus', ' a photo of a marigold', ' a photo of a marigold', ' a photo of a rose', ' a photo of a rose', ' a photo of a rose',\n",
    "               ' orange dahlia', ' orange dahlia', ' a lenten rose', ' a lenten rose', ' a water lily', ' a water lily', ' a water lily', ' a water lily', ' an orchid', ' an orchid', ' an orchid', ' hard-leaved pocket orchid', ' bird of paradise', ' bird of paradise', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a globe-flower', ' a photo of a globe-flower', ' a photo of a lovely rose', ' a photo of a lovely rose', ' a photo of a ruby-lipped cattleya', ' a photo of a ruby-lipped cattleya', ' a photo of a lovely rose', ' a water lily', ' a osteospermum', ' a osteospermum', ' a water lily', ' a water lily', ' a water lily', ' a red rose', ' a red rose']\n",
    "\n",
    "\n",
    "def get_val_dataset(batch_size=8):\n",
    "    for i in range(0, len(val_prompts), batch_size):\n",
    "        prompts = val_prompts[i:i + batch_size]\n",
    "        tokens = text_encoder.tokenize(prompts)\n",
    "        yield {\"text\": tokens}\n",
    "\n",
    "\n",
    "data['test'] = get_val_dataset\n",
    "data['test_len'] = len(val_prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculated input shapes: {'x': (32, 32, 4), 'temb': (), 'textcontext': (77, 768)}\n"
     ]
    }
   ],
   "source": [
    "from flax import linen as nn\n",
    "from diffusers import FlaxUNet2DConditionModel\n",
    "from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig\n",
    "\n",
    "input_config = DiffusionInputConfig(\n",
    "    sample_data_key='image',\n",
    "    sample_data_shape=(IMAGE_SIZE, IMAGE_SIZE, 3),\n",
    "    conditions=[\n",
    "        ConditionalInputConfig(\n",
    "            encoder=text_encoder,\n",
    "            conditioning_data_key='text',\n",
    "            pretokenized=True,\n",
    "            unconditional_input=\"\",\n",
    "            model_key_override=\"textcontext\",\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "\n",
    "input_shapes = input_config.get_input_shapes(\n",
    "    autoencoder=autoencoder,\n",
    ")\n",
    "\n",
    "unet_model = FlaxUNet2DConditionModel(\n",
    "    sample_size=input_shapes[\"x\"][1],  # the target image resolution\n",
    "    # the number of input channels, 3 for RGB images\n",
    "    in_channels=input_shapes[\"x\"][2],\n",
    "    out_channels=input_shapes[\"x\"][2],  # the number of output channels\n",
    "    layers_per_block=2,  # how many ResNet layers to use per UNet block\n",
    "    # the number of output channels for each UNet block\n",
    "    block_out_channels=(64, 128, 256, 512),\n",
    "    cross_attention_dim=512,  # the size of the cross-attention layers\n",
    "    dtype=jnp.bfloat16,\n",
    "    use_memory_efficient_attention=True,\n",
    ")\n",
    "\n",
    "\n",
    "class BCHWModelWrapper(nn.Module):\n",
    "    model: nn.Module\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x, temb, textcontext):\n",
    "        # Reshape the input to BCHW format from BHWC\n",
    "        x = jnp.transpose(x, (0, 3, 1, 2))\n",
    "        # Pass the input through the UNet model\n",
    "        out = self.model(\n",
    "            sample=x,\n",
    "            timesteps=temb,\n",
    "            encoder_hidden_states=textcontext,\n",
    "        )\n",
    "        # Reshape the output back to BHWC format\n",
    "        out = jnp.transpose(out.sample, (0, 2, 3, 1))\n",
    "        return out\n",
    "    \n",
    "    @property\n",
    "    def __dict__(self):\n",
    "        return self.model.__dict__\n",
    "\n",
    "unet = BCHWModelWrapper(unet_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculated input shapes: {'x': (32, 32, 4), 'temb': (), 'textcontext': (77, 768)}\n"
     ]
    }
   ],
   "source": [
    "from flaxdiff.inputs import DiffusionInputConfig, ConditionalInputConfig\n",
    "\n",
    "input_config = DiffusionInputConfig(\n",
    "    sample_data_key='image',\n",
    "    sample_data_shape=(IMAGE_SIZE, IMAGE_SIZE, 3),\n",
    "    conditions=[\n",
    "        ConditionalInputConfig(\n",
    "            encoder=text_encoder,\n",
    "            conditioning_data_key='text',\n",
    "            pretokenized=True,\n",
    "            unconditional_input=\"\",\n",
    "            model_key_override=\"textcontext\",\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "input_shapes = input_config.get_input_shapes(\n",
    "    autoencoder=autoencoder,\n",
    ")\n",
    "\n",
    "unet = Unet(emb_features=256,\n",
    "            feature_depths=[64, 64, 128, 256, 512],\n",
    "            attention_configs=[\n",
    "                None,\n",
    "                {\"heads\": 8, \"dtype\": jnp.float32, \"flash_attention\": False,\n",
    "                    \"use_projection\": False, \"use_self_and_cross\": True},\n",
    "                {\"heads\": 8, \"dtype\": jnp.float32, \"flash_attention\": False,\n",
    "                    \"use_projection\": False, \"use_self_and_cross\": True},\n",
    "                {\"heads\": 8, \"dtype\": jnp.float32, \"flash_attention\": False,\n",
    "                    \"use_projection\": False, \"use_self_and_cross\": True},\n",
    "                {\"heads\": 8, \"dtype\": jnp.float32, \"flash_attention\": False,\n",
    "                    \"use_projection\": False, \"use_self_and_cross\": False}\n",
    "            ],\n",
    "            num_res_blocks=2,\n",
    "            num_middle_res_blocks=1,\n",
    "            dtype=jnp.bfloat16,\n",
    "            output_channels=input_shapes[\"x\"][2],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculated input shapes: {'x': (32, 32, 4), 'temb': (), 'textcontext': (77, 768)}\n",
      "Model name: diffusion-oxford_flowers102-res256\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mashishkumar4\u001b[0m to \u001b[32mhttps://api.wandb.ai\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.19.9"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/mrwhite0racle/persist/FlaxDiff/wandb/run-20250419_153720-4b2mer47</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/umd-projects/mlops-msml605-project/runs/4b2mer47' target=\"_blank\">General_Diffusion_demo_for_inference2</a></strong> to <a href='https://wandb.ai/umd-projects/mlops-msml605-project' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/umd-projects/mlops-msml605-project' target=\"_blank\">https://wandb.ai/umd-projects/mlops-msml605-project</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/umd-projects/mlops-msml605-project/runs/4b2mer47' target=\"_blank\">https://wandb.ai/umd-projects/mlops-msml605-project/runs/4b2mer47</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate.\n",
      "WARNING:absl:Configured `CheckpointManager` using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from checkpoint at step  411355\n",
      "Loaded model from checkpoint at epoch 804 step 411355 0.4702588\n",
      "Generating states for DiffusionTrainer\n"
     ]
    }
   ],
   "source": [
    "# Define noise scheduler\n",
    "edm_schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)\n",
    "karas_ve_schedule = KarrasVENoiseScheduler(\n",
    "    1, sigma_max=80, rho=7, sigma_data=0.5)\n",
    "# Define model\n",
    "\n",
    "# Define optimizer\n",
    "solver = optax.adam(2e-4)\n",
    "\n",
    "# Create the GeneralDiffusionTrainer\n",
    "experiment_name = \"General_Diffusion_2025-04-18_06:34:50\"#f\"General_Diffusion_{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}\"\n",
    "\n",
    "trainer = GeneralDiffusionTrainer(\n",
    "    unet,\n",
    "    optimizer=solver,\n",
    "    noise_schedule=edm_schedule,\n",
    "    autoencoder=autoencoder,\n",
    "    input_config=input_config,\n",
    "    rngs=jax.random.PRNGKey(42),\n",
    "    name=experiment_name,\n",
    "    model_output_transform=KarrasPredictionTransform(\n",
    "        sigma_data=edm_schedule.sigma_data),\n",
    "    # data_key='image',  # Specify the key for image data in batches\n",
    "    distributed_training=True,\n",
    "    wandb_config={\n",
    "        \"project\": 'mlops-msml605-project',\n",
    "        \"entity\": 'umd-projects',\n",
    "        \"name\": experiment_name,\n",
    "        \"id\": \"bdw4ebqf\",\n",
    "        \"config\": {\n",
    "            \"batch_size\": BATCH_SIZE,\n",
    "            \"image_size\": IMAGE_SIZE,\n",
    "            \"arguments\": {\n",
    "                \"architecture\": \"unet\",\n",
    "                \"dataset\": \"oxford_flowers102\",\n",
    "                \"noise_schedule\": \"edm\",\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    native_resolution=IMAGE_SIZE,\n",
    "    # Path to the checkpoint\n",
    "    load_from_checkpoint=\"/home/mrwhite0racle/persist/FlaxDiff/checkpoints/general_diffusion_2025-04-18_06:34:50\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 897: 600step [00:31, 19.01step/s, loss=0.4680]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 897 Loss: 0.47235533595085144\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 897 completed. Avg Loss: 0.47235533595085144, Time: 31.57s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 105.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 898/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 898:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.4724]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 458878\n",
      "Training started for process index 0 at step 458878\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 898: 600step [00:30, 19.66step/s, loss=0.4214]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 898 Loss: 0.4722280502319336\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 898 completed. Avg Loss: 0.4722280502319336, Time: 30.52s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 104.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 899/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 899:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.4720]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 459389\n",
      "Training started for process index 0 at step 459389\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 899: 600step [00:30, 19.51step/s, loss=0.5028]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 899 Loss: 0.473550945520401\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 899 completed. Avg Loss: 0.473550945520401, Time: 30.75s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 104.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 900/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 900:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.5043]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 459900\n",
      "Training started for process index 0 at step 459900\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 900: 600step [00:30, 19.69step/s, loss=0.4644]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 900 Loss: 0.47377240657806396\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 900 completed. Avg Loss: 0.47377240657806396, Time: 30.48s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 101.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 901/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 901:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.4685]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 460411\n",
      "Training started for process index 0 at step 460411\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 901: 600step [00:31, 19.08step/s, loss=0.4860]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 901 Loss: 0.47391751408576965\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 901 completed. Avg Loss: 0.47391751408576965, Time: 31.45s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 100.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 902/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 902:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.4568]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 460922\n",
      "Training started for process index 0 at step 460922\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 902: 600step [00:31, 18.90step/s, loss=0.4688]                                              "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mEpoch done on index 0 => 902 Loss: 0.47201600670814514\u001b[0m\n",
      "\u001b[32mEpoch done on process index 0\u001b[0m\n",
      "\u001b[32m\n",
      "\tEpoch 902 completed. Avg Loss: 0.47201600670814514, Time: 31.75s, Best Loss: 0.47025880217552185\u001b[0m\n",
      "Validation started for process index 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "100%|██████████| 200/200 [00:01<00:00, 100.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32mValidation done on process index 0\u001b[0m\n",
      "\n",
      "Epoch 903/2000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\t\tEpoch 903:   0%|                                           | 0/511 [00:00<?, ?step/s, loss=0.4279]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First batch loaded at step 461433\n",
      "Training started for process index 0 at step 461433\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Process SpawnProcess-29:\n",
      "Process SpawnProcess-32:\n",
      "Process SpawnProcess-15:\n",
      "Process SpawnProcess-18:\n",
      "Process SpawnProcess-10:\n",
      "Process SpawnProcess-31:\n",
      "Process SpawnProcess-21:\n",
      "Process SpawnProcess-6:\n",
      "Process SpawnProcess-16:\n",
      "Process SpawnProcess-27:\n",
      "Process SpawnProcess-9:\n",
      "Process SpawnProcess-4:\n",
      "Process SpawnProcess-25:\n",
      "Process SpawnProcess-7:\n",
      "Process SpawnProcess-23:\n",
      "Process SpawnProcess-11:\n",
      "Process SpawnProcess-28:\n",
      "Process SpawnProcess-12:\n",
      "Process SpawnProcess-30:\n",
      "Process SpawnProcess-19:\n",
      "Process SpawnProcess-2:\n",
      "Process SpawnProcess-5:\n",
      "Process SpawnProcess-8:\n",
      "Process SpawnProcess-14:\n",
      "Process SpawnProcess-24:\n",
      "Process SpawnProcess-26:\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Process SpawnProcess-3:\n",
      "Process SpawnProcess-20:\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "Process SpawnProcess-22:\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Process SpawnProcess-17:\n",
      "Process SpawnProcess-13:\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "KeyboardInterrupt\n",
      "Process SpawnProcess-1:\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 235, in _worker_loop\n",
      "    next_element = next(element_producer)\n",
      "                   ^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/data_loader.py\", line 387, in __call__\n",
      "    yield from self._read_and_transform_data(last_seen_index)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/data_loader.py\", line 528, in _apply_transform\n",
      "    for input_record in input_iterator:\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/data_loader.py\", line 518, in _apply_transform\n",
      "    for r in batch_op(input_iterator):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/operations.py\", line 152, in __call__\n",
      "    for input_record in input_iterator:\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/data_loader.py\", line 530, in _apply_transform\n",
      "    output_record, filter_result = fn(input_record)\n",
      "                                   ^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/data_loader.py\", line 497, in <lambda>\n",
      "    fn = lambda r: (record.Record(r.metadata, transform.map(r.data)), True)\n",
      "                                              ^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/persist/FlaxDiff/flaxdiff/data/sources/images.py\", line 159, in map\n",
      "    results = self.tokenize(caption)\n",
      "              ^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/persist/FlaxDiff/flaxdiff/utils.py\", line 166, in __call__\n",
      "    tokens = self.tokenizer(inputs, padding=\"max_length\", max_length=self.tokenizer.model_max_length,\n",
      "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/transformers/tokenization_utils_base.py\", line 2887, in __call__\n",
      "    encodings = self._call_one(text=text, text_pair=text_pair, **all_kwargs)\n",
      "                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/transformers/tokenization_utils_base.py\", line 2997, in _call_one\n",
      "    return self.encode_plus(\n",
      "           ^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/transformers/tokenization_utils_base.py\", line 3073, in encode_plus\n",
      "    return self._encode_plus(\n",
      "           ^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/transformers/tokenization_utils_fast.py\", line 613, in _encode_plus\n",
      "    batched_output = self._batch_encode_plus(\n",
      "                     ^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/transformers/tokenization_utils_fast.py\", line 539, in _batch_encode_plus\n",
      "    encodings = self._tokenizer.encode_batch(\n",
      "                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "KeyboardInterrupt\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n",
      "Traceback (most recent call last):\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 314, in _bootstrap\n",
      "    self.run()\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/process.py\", line 108, in run\n",
      "    self._target(*self._args, **self._kwargs)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/grain_pool.py\", line 236, in _worker_loop\n",
      "    if not multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/grain/_src/python/multiprocessing_common.py\", line 54, in add_element_to_queue\n",
      "    elements_queue.put(element, timeout=_QUEUE_WAIT_TIMEOUT_SECONDS)\n",
      "  File \"/home/mrwhite0racle/miniconda3/envs/flaxdiff/lib/python3.11/multiprocessing/queues.py\", line 89, in put\n",
      "    if not self._sem.acquire(block, timeout):\n",
      "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "KeyboardInterrupt\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 2\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;66;03m# Train the model\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m2\u001b[39m final_state = \u001b[43mtrainer\u001b[49m\u001b[43m.\u001b[49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatches\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m2000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msampler_class\u001b[49m\u001b[43m=\u001b[49m\u001b[43mEulerAncestralSampler\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msampling_noise_schedule\u001b[49m\u001b[43m=\u001b[49m\u001b[43mkaras_ve_schedule\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/persist/FlaxDiff/flaxdiff/trainer/diffusion_trainer.py:361\u001b[39m, in \u001b[36mDiffusionTrainer.fit\u001b[39m\u001b[34m(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch, sampler_class, sampling_noise_schedule)\u001b[39m\n\u001b[32m    356\u001b[39m local_batch_size = data[\u001b[33m'\u001b[39m\u001b[33mlocal_batch_size\u001b[39m\u001b[33m'\u001b[39m]\n\u001b[32m    357\u001b[39m validation_step_args = {\n\u001b[32m    358\u001b[39m     \u001b[33m\"\u001b[39m\u001b[33msampler_class\u001b[39m\u001b[33m\"\u001b[39m: sampler_class,\n\u001b[32m    359\u001b[39m     \u001b[33m\"\u001b[39m\u001b[33msampling_noise_schedule\u001b[39m\u001b[33m\"\u001b[39m: sampling_noise_schedule,\n\u001b[32m    360\u001b[39m }\n\u001b[32m--> \u001b[39m\u001b[32m361\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m.\u001b[49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    362\u001b[39m \u001b[43m    \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m    363\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_steps_per_epoch\u001b[49m\u001b[43m=\u001b[49m\u001b[43mtraining_steps_per_epoch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m    364\u001b[39m \u001b[43m    \u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m=\u001b[49m\u001b[43mepochs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m    365\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_step_args\u001b[49m\u001b[43m=\u001b[49m\u001b[43m{\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mbatch_size\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mlocal_batch_size\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[32m    366\u001b[39m \u001b[43m    \u001b[49m\u001b[43mval_steps_per_epoch\u001b[49m\u001b[43m=\u001b[49m\u001b[43mval_steps_per_epoch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    367\u001b[39m \u001b[43m    \u001b[49m\u001b[43mvalidation_step_args\u001b[49m\u001b[43m=\u001b[49m\u001b[43mvalidation_step_args\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    368\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/persist/FlaxDiff/flaxdiff/trainer/simple_trainer.py:526\u001b[39m, in \u001b[36mSimpleTrainer.fit\u001b[39m\u001b[34m(self, data, train_steps_per_epoch, epochs, train_step_args, val_steps_per_epoch, validation_step_args)\u001b[39m\n\u001b[32m    523\u001b[39m start_time = time.time()\n\u001b[32m    524\u001b[39m epoch_loss = \u001b[32m0\u001b[39m\n\u001b[32m--> \u001b[39m\u001b[32m526\u001b[39m epoch_loss, current_step, train_state, rng_state = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mtrain_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    527\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_state\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    528\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_step\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    529\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_ds\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    530\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtrain_steps_per_epoch\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    531\u001b[39m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mlatest_step\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    532\u001b[39m \u001b[43m    \u001b[49m\u001b[43mrng_state\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    533\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    534\u001b[39m \u001b[38;5;28mprint\u001b[39m(colored(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mEpoch done on process index \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprocess_index\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m, PROCESS_COLOR_MAP[process_index]))\n\u001b[32m    536\u001b[39m \u001b[38;5;28mself\u001b[39m.latest_step = current_step\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/persist/FlaxDiff/flaxdiff/trainer/simple_trainer.py:438\u001b[39m, in \u001b[36mSimpleTrainer.train_loop\u001b[39m\u001b[34m(self, train_state, train_step_fn, train_ds, train_steps_per_epoch, current_step, rng_state)\u001b[39m\n\u001b[32m    434\u001b[39m     \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mTraining started for process index \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mprocess_index\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m at step \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcurrent_step\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m)\n\u001b[32m    436\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.distributed_training:\n\u001b[32m    437\u001b[39m     \u001b[38;5;66;03m# loss = jax.experimental.multihost_utils.process_allgather(loss)\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m438\u001b[39m     loss = \u001b[43mjnp\u001b[49m\u001b[43m.\u001b[49m\u001b[43mmean\u001b[49m\u001b[43m(\u001b[49m\u001b[43mloss\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;66;03m# Just to make sure its a scaler value\u001b[39;00m\n\u001b[32m    440\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m loss <= \u001b[32m1e-8\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m jnp.isnan(loss) \u001b[38;5;129;01mor\u001b[39;00m jnp.isinf(loss):\n\u001b[32m    441\u001b[39m     \u001b[38;5;66;03m# If the loss is too low or NaN/Inf, log the issue and attempt recovery\u001b[39;00m\n\u001b[32m    442\u001b[39m     \u001b[38;5;28mprint\u001b[39m(colored(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mAbnormal loss at step \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mcurrent_step\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mloss\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m\"\u001b[39m, \u001b[33m'\u001b[39m\u001b[33mred\u001b[39m\u001b[33m'\u001b[39m))\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/jax/_src/numpy/reductions.py:803\u001b[39m, in \u001b[36mmean\u001b[39m\u001b[34m(a, axis, dtype, out, keepdims, where)\u001b[39m\n\u001b[32m    799\u001b[39m     size *= maybe_named_axis(a, \u001b[38;5;28;01mlambda\u001b[39;00m i: a_shape[i], \u001b[38;5;28;01mlambda\u001b[39;00m name: lax.psum(\u001b[32m1\u001b[39m, name))\n\u001b[32m    800\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m size\n\u001b[32m--> \u001b[39m\u001b[32m803\u001b[39m \u001b[38;5;129m@export\u001b[39m\n\u001b[32m    804\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mmean\u001b[39m(a: ArrayLike, axis: Axis = \u001b[38;5;28;01mNone\u001b[39;00m, dtype: DTypeLike | \u001b[38;5;28;01mNone\u001b[39;00m = \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[32m    805\u001b[39m          out: \u001b[38;5;28;01mNone\u001b[39;00m = \u001b[38;5;28;01mNone\u001b[39;00m, keepdims: \u001b[38;5;28mbool\u001b[39m = \u001b[38;5;28;01mFalse\u001b[39;00m, *,\n\u001b[32m    806\u001b[39m          where: ArrayLike | \u001b[38;5;28;01mNone\u001b[39;00m = \u001b[38;5;28;01mNone\u001b[39;00m) -> Array:\n\u001b[32m    807\u001b[39m \u001b[38;5;250m  \u001b[39m\u001b[33mr\u001b[39m\u001b[33;03m\"\"\"Return the mean of array elements along a given axis.\u001b[39;00m\n\u001b[32m    808\u001b[39m \n\u001b[32m    809\u001b[39m \u001b[33;03m  JAX implementation of :func:`numpy.mean`.\u001b[39;00m\n\u001b[32m   (...)\u001b[39m\u001b[32m    865\u001b[39m \u001b[33;03m           [6. ]], dtype=float32)\u001b[39;00m\n\u001b[32m    866\u001b[39m \u001b[33;03m  \"\"\"\u001b[39;00m\n\u001b[32m    867\u001b[39m   \u001b[38;5;28;01mreturn\u001b[39;00m _mean(a, _ensure_optional_axes(axis), dtype, out, keepdims,\n\u001b[32m    868\u001b[39m                where=where, upcast_f16_for_computation=(dtype \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m))\n",
      "\u001b[31mKeyboardInterrupt\u001b[39m: "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error in callback <bound method _WandbInit._pause_backend of <wandb.sdk.wandb_init._WandbInit object at 0x7f734c1c4290>> (for post_run_cell), with arguments args (<ExecutionResult object at 7f726c109890, execution_count=5 error_before_exec=None error_in_exec= info=<ExecutionInfo object at 7f70ac1bd0d0, raw_cell=\"# Train the model\n",
      "final_state = trainer.fit(data, ..\" store_history=True silent=False shell_futures=True cell_id=vscode-notebook-cell://ssh-remote%2Btpu-v4-8/home/mrwhite0racle/persist/FlaxDiff/prototype_general_pipeline.ipynb#W5sdnNjb2RlLXJlbW90ZQ%3D%3D> result=None>,),kwargs {}:\n"
     ]
    },
    {
     "ename": "MailboxClosedError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mMailboxClosedError\u001b[39m                        Traceback (most recent call last)",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_init.py:543\u001b[39m, in \u001b[36m_WandbInit._pause_backend\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m    541\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.notebook.save_ipynb():  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[32m    542\u001b[39m     \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m.run \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m543\u001b[39m     res = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mrun\u001b[49m\u001b[43m.\u001b[49m\u001b[43mlog_code\u001b[49m\u001b[43m(\u001b[49m\u001b[43mroot\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[32m    544\u001b[39m     \u001b[38;5;28mself\u001b[39m._logger.info(\u001b[33m\"\u001b[39m\u001b[33msaved code: \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[33m\"\u001b[39m, res)  \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[32m    545\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m.backend.interface \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_run.py:435\u001b[39m, in \u001b[36m_run_decorator._noop_on_finish.<locals>.decorator_fn.<locals>.wrapper_fn\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m    432\u001b[39m \u001b[38;5;129m@functools\u001b[39m.wraps(func)\n\u001b[32m    433\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mwrapper_fn\u001b[39m(\u001b[38;5;28mself\u001b[39m: \u001b[38;5;28mtype\u001b[39m[Run], *args: Any, **kwargs: Any) -> Any:\n\u001b[32m    434\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33m_is_finished\u001b[39m\u001b[33m\"\u001b[39m, \u001b[38;5;28;01mFalse\u001b[39;00m):\n\u001b[32m--> \u001b[39m\u001b[32m435\u001b[39m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    437\u001b[39m     default_message = (\n\u001b[32m    438\u001b[39m         \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mRun (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m.id\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m) is finished. The call to `\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfunc.\u001b[34m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m` will be ignored. \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    439\u001b[39m         \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mPlease make sure that you are using an active run.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    440\u001b[39m     )\n\u001b[32m    441\u001b[39m     resolved_message = message \u001b[38;5;129;01mor\u001b[39;00m default_message\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_run.py:387\u001b[39m, in \u001b[36m_log_to_run.<locals>.wrapper\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m    384\u001b[39m     run_id = \u001b[38;5;28mself\u001b[39m._attach_id\n\u001b[32m    386\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m wb_logging.log_to_run(run_id):\n\u001b[32m--> \u001b[39m\u001b[32m387\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_run.py:425\u001b[39m, in \u001b[36m_run_decorator._attach.<locals>.wrapper\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m    423\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m e\n\u001b[32m    424\u001b[39m     \u001b[38;5;28mcls\u001b[39m._is_attaching = \u001b[33m\"\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m--> \u001b[39m\u001b[32m425\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_run.py:1147\u001b[39m, in \u001b[36mRun.log_code\u001b[39m\u001b[34m(self, root, name, include_fn, exclude_fn)\u001b[39m\n\u001b[32m   1142\u001b[39m     wandb.termwarn(\n\u001b[32m   1143\u001b[39m         \u001b[33m\"\u001b[39m\u001b[33mNo relevant files were detected in the specified directory. No code will be logged to your run.\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m   1144\u001b[39m     )\n\u001b[32m   1145\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1147\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_log_artifact\u001b[49m\u001b[43m(\u001b[49m\u001b[43mart\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/wandb_run.py:3351\u001b[39m, in \u001b[36mRun._log_artifact\u001b[39m\u001b[34m(self, artifact_or_path, name, type, aliases, tags, distributed_id, finalize, is_user_created, use_after_commit)\u001b[39m\n\u001b[32m   3349\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backend \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backend.interface:\n\u001b[32m   3350\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m._settings._offline:\n\u001b[32m-> \u001b[39m\u001b[32m3351\u001b[39m         handle = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_backend\u001b[49m\u001b[43m.\u001b[49m\u001b[43minterface\u001b[49m\u001b[43m.\u001b[49m\u001b[43mdeliver_artifact\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m   3352\u001b[39m \u001b[43m            \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m   3353\u001b[39m \u001b[43m            \u001b[49m\u001b[43martifact\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3354\u001b[39m \u001b[43m            \u001b[49m\u001b[43maliases\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3355\u001b[39m \u001b[43m            \u001b[49m\u001b[43mtags\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3356\u001b[39m \u001b[43m            \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mstep\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3357\u001b[39m \u001b[43m            \u001b[49m\u001b[43mfinalize\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfinalize\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3358\u001b[39m \u001b[43m            \u001b[49m\u001b[43mis_user_created\u001b[49m\u001b[43m=\u001b[49m\u001b[43mis_user_created\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3359\u001b[39m \u001b[43m            \u001b[49m\u001b[43muse_after_commit\u001b[49m\u001b[43m=\u001b[49m\u001b[43muse_after_commit\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m   3360\u001b[39m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   3361\u001b[39m         artifact._set_save_handle(handle, \u001b[38;5;28mself\u001b[39m._public_api().client)\n\u001b[32m   3362\u001b[39m     \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/interface/interface.py:589\u001b[39m, in \u001b[36mInterfaceBase.deliver_artifact\u001b[39m\u001b[34m(self, run, artifact, aliases, tags, history_step, is_user_created, use_after_commit, finalize)\u001b[39m\n\u001b[32m    587\u001b[39m     log_artifact.history_step = history_step\n\u001b[32m    588\u001b[39m log_artifact.staging_dir = get_staging_dir()\n\u001b[32m--> \u001b[39m\u001b[32m589\u001b[39m resp = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_deliver_artifact\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlog_artifact\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    590\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m resp\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/interface/interface_shared.py:339\u001b[39m, in \u001b[36mInterfaceShared._deliver_artifact\u001b[39m\u001b[34m(self, log_artifact)\u001b[39m\n\u001b[32m    334\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_deliver_artifact\u001b[39m(\n\u001b[32m    335\u001b[39m     \u001b[38;5;28mself\u001b[39m,\n\u001b[32m    336\u001b[39m     log_artifact: pb.LogArtifactRequest,\n\u001b[32m    337\u001b[39m ) -> MailboxHandle[pb.Result]:\n\u001b[32m    338\u001b[39m     rec = \u001b[38;5;28mself\u001b[39m._make_request(log_artifact=log_artifact)\n\u001b[32m--> \u001b[39m\u001b[32m339\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_deliver_record\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrec\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/interface/interface_shared.py:389\u001b[39m, in \u001b[36mInterfaceShared._deliver_record\u001b[39m\u001b[34m(self, record)\u001b[39m\n\u001b[32m    386\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34m_deliver_record\u001b[39m(\u001b[38;5;28mself\u001b[39m, record: pb.Record) -> MailboxHandle[pb.Result]:\n\u001b[32m    387\u001b[39m     mailbox = \u001b[38;5;28mself\u001b[39m._get_mailbox()\n\u001b[32m--> \u001b[39m\u001b[32m389\u001b[39m     handle = \u001b[43mmailbox\u001b[49m\u001b[43m.\u001b[49m\u001b[43mrequire_response\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrecord\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    390\u001b[39m     \u001b[38;5;28mself\u001b[39m._publish(record)\n\u001b[32m    392\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m handle.map(\u001b[38;5;28;01mlambda\u001b[39;00m resp: resp.result_communicate)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/miniconda3/envs/flaxdiff/lib/python3.11/site-packages/wandb/sdk/mailbox/mailbox.py:68\u001b[39m, in \u001b[36mMailbox.require_response\u001b[39m\u001b[34m(self, request)\u001b[39m\n\u001b[32m     66\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m._handles_lock:\n\u001b[32m     67\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m._closed:\n\u001b[32m---> \u001b[39m\u001b[32m68\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m MailboxClosedError()\n\u001b[32m     70\u001b[39m     handle = MailboxResponseHandle(address)\n\u001b[32m     71\u001b[39m     \u001b[38;5;28mself\u001b[39m._handles[address] = handle\n",
      "\u001b[31mMailboxClosedError\u001b[39m: "
     ]
    }
   ],
   "source": [
    "# Train the model\n",
    "final_state = trainer.fit(data, batches, epochs=2000,\n",
    "                          sampler_class=EulerAncestralSampler, sampling_noise_schedule=karas_ve_schedule)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def normalizeImage(x): return jax.nn.standardize(x, mean=[127.5], std=[127.5])\n",
    "def denormalizeImage(x): return (x + 1.0) * 127.5\n",
    "\n",
    "\n",
    "def plotImages(imgs, fig_size=(8, 8), dpi=100):\n",
    "    fig = plt.figure(figsize=fig_size, dpi=dpi)\n",
    "    imglen = imgs.shape[0]\n",
    "    for i in range(imglen):\n",
    "        plt.subplot(fig_size[0], fig_size[1], i + 1)\n",
    "        plt.imshow(jnp.astype(denormalizeImage(imgs[i, :, :, :]), jnp.uint8))\n",
    "        plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using classifier-free guidance\n"
     ]
    }
   ],
   "source": [
    "sampler = EulerAncestralSampler(\n",
    "    model=trainer.model,\n",
    "    noise_schedule=karas_ve_schedule,\n",
    "    model_output_transform=KarrasPredictionTransform(\n",
    "        sigma_data=karas_ve_schedule.sigma_data),\n",
    "    autoencoder=trainer.autoencoder,\n",
    "    input_config=trainer.input_config,\n",
    "    guidance_scale=3,\n",
    "    timestep_spacing=\"linear\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing raw conditioning inputs to generate model conditioning inputs\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:15<00:00, 13.01it/s]\n"
     ]
    }
   ],
   "source": [
    "prompts = [\n",
    "    'water tulip',\n",
    "    'a water lily',\n",
    "    'a water lily',\n",
    "    'a photo of a rose',\n",
    "    'a photo of a rose',\n",
    "    'a water lily',\n",
    "    'a water lily',\n",
    "    'a photo of a marigold',\n",
    "]\n",
    "images = sampler.generate_samples(\n",
    "    params=trainer.best_state.params,\n",
    "    resolution=IMAGE_SIZE,\n",
    "    num_samples=len(prompts),\n",
    "    sequence_length=None,\n",
    "    diffusion_steps=200,\n",
    "    start_step=1000,\n",
    "    end_step=0,\n",
    "    conditioning=prompts,\n",
    "    # model_conditioning_inputs=(encoded,)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'plotImages' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[5]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mplotImages\u001b[49m(images, dpi=\u001b[32m500\u001b[39m)\n",
      "\u001b[31mNameError\u001b[39m: name 'plotImages' is not defined"
     ]
    }
   ],
   "source": [
    "plotImages(images, dpi=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/home/mrwhite0racle/persist/FlaxDiff/checkpoints/general_diffusion_demo_for_inference'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.checkpoint_path()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Adding directory to artifact (/home/mrwhite0racle/persist/FlaxDiff/checkpoints/general_diffusion_demo_for_inference2/411355)... Done. 8.9s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model pushed to registry at wandb-registry-model/diffusion-oxford_flowers102-res256\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Artifact QXJ0aWZhY3Q6MTY2NzkxMDU2OA==>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.push_to_registry()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'trainer' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[43mtrainer\u001b[49m.wandb.run\n",
      "\u001b[31mNameError\u001b[39m: name 'trainer' is not defined"
     ]
    }
   ],
   "source": [
    "trainer.wandb.run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "flaxdiff",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}