Upload folder using huggingface_hub

.ipynb_checkpoints/pipeline_waifu-checkpoint.py

@@ -1,5 +1,6 @@
 import torch
 from diffusers import DiffusionPipeline
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 # waifu
 # tokenizer

Untitled.ipynb

@@ -2,28 +2,14 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 16,
    "id": "dca3239c-17d6-4284-a2cf-83237a55a7df",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2e20307c4a1b4f2981e89087b8613a78",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "model_index.json:   0%|          | 0.00/475 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7c51603173654d1eb786d943b0489ea7",
+       "model_id": "0314d609c09a4ad3b8c9d432e549330f",
        "version_major": 2,
        "version_minor": 0
       },
@@ -36,12 +22,12 @@
     },
     {
      "ename": "NameError",
-     "evalue": "name 'DiffusionPipeline' is not defined",
+     "evalue": "name 'Union' is not defined",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[14], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m pipe_id \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAiArtLab/waifu-2b\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      5\u001b[0m variant \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfp16\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 6\u001b[0m pipe \u001b[38;5;241m=\u001b[39m \u001b[43mDiffusionPipeline\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpipe_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43mvariant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvariant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrust_remote_code\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\n\u001b[1;32m     10\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(pipe)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m#pipe_sd.to(\"cuda\")\u001b[39;00m\n",
+      "Cell \u001b[0;32mIn[16], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m pipe_id \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAiArtLab/waifu-2b\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      5\u001b[0m variant \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfp16\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 6\u001b[0m pipe \u001b[38;5;241m=\u001b[39m \u001b[43mDiffusionPipeline\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43mpipe_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43mvariant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvariant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrust_remote_code\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\n\u001b[1;32m     10\u001b[0m \u001b[43m)\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(pipe)\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m#pipe_sd.to(\"cuda\")\u001b[39;00m\n",
       "File \u001b[0;32m~/.local/lib/python3.11/site-packages/huggingface_hub/utils/_validators.py:114\u001b[0m, in \u001b[0;36mvalidate_hf_hub_args.<locals>._inner_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m check_use_auth_token:\n\u001b[1;32m    112\u001b[0m     kwargs \u001b[38;5;241m=\u001b[39m smoothly_deprecate_use_auth_token(fn_name\u001b[38;5;241m=\u001b[39mfn\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, has_token\u001b[38;5;241m=\u001b[39mhas_token, kwargs\u001b[38;5;241m=\u001b[39mkwargs)\n\u001b[0;32m--> 114\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
       "File \u001b[0;32m~/.local/lib/python3.11/site-packages/diffusers/pipelines/pipeline_utils.py:731\u001b[0m, in \u001b[0;36mDiffusionPipeline.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, **kwargs)\u001b[0m\n\u001b[1;32m    726\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m pretrained_model_name_or_path\u001b[38;5;241m.\u001b[39mcount(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/\u001b[39m\u001b[38;5;124m\"\u001b[39m) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    727\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    728\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mThe provided pretrained_model_name_or_path \u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpretrained_model_name_or_path\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m    729\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m is neither a valid local path nor a valid repo id. Please check the parameter.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    730\u001b[0m         )\n\u001b[0;32m--> 731\u001b[0m     cached_folder \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdownload\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    732\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    733\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcache_dir\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_dir\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    734\u001b[0m \u001b[43m        \u001b[49m\u001b[43mforce_download\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mforce_download\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    735\u001b[0m \u001b[43m        \u001b[49m\u001b[43mproxies\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mproxies\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    736\u001b[0m \u001b[43m        \u001b[49m\u001b[43mlocal_files_only\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlocal_files_only\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    737\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtoken\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    738\u001b[0m \u001b[43m        \u001b[49m\u001b[43mrevision\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrevision\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    739\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfrom_flax\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfrom_flax\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    740\u001b[0m \u001b[43m        \u001b[49m\u001b[43muse_safetensors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43muse_safetensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    741\u001b[0m \u001b[43m        \u001b[49m\u001b[43muse_onnx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43muse_onnx\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    742\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcustom_pipeline\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcustom_pipeline\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    743\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcustom_revision\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcustom_revision\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    744\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvariant\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvariant\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    745\u001b[0m \u001b[43m        \u001b[49m\u001b[43mload_connected_pipeline\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mload_connected_pipeline\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    746\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    747\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    748\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    749\u001b[0m     cached_folder \u001b[38;5;241m=\u001b[39m pretrained_model_name_or_path\n",
       "File \u001b[0;32m~/.local/lib/python3.11/site-packages/huggingface_hub/utils/_validators.py:114\u001b[0m, in \u001b[0;36mvalidate_hf_hub_args.<locals>._inner_fn\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m check_use_auth_token:\n\u001b[1;32m    112\u001b[0m     kwargs \u001b[38;5;241m=\u001b[39m smoothly_deprecate_use_auth_token(fn_name\u001b[38;5;241m=\u001b[39mfn\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m, has_token\u001b[38;5;241m=\u001b[39mhas_token, kwargs\u001b[38;5;241m=\u001b[39mkwargs)\n\u001b[0;32m--> 114\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
@@ -58,8 +44,9 @@
       "File \u001b[0;32m<frozen importlib._bootstrap>:690\u001b[0m, in \u001b[0;36m_load_unlocked\u001b[0;34m(spec)\u001b[0m\n",
       "File \u001b[0;32m<frozen importlib._bootstrap_external>:940\u001b[0m, in \u001b[0;36mexec_module\u001b[0;34m(self, module)\u001b[0m\n",
       "File \u001b[0;32m<frozen importlib._bootstrap>:241\u001b[0m, in \u001b[0;36m_call_with_frames_removed\u001b[0;34m(f, *args, **kwds)\u001b[0m\n",
-      "File \u001b[0;32m~/.cache/huggingface/modules/diffusers_modules/local/AiArtLab--waifu-2b/3d54a9ff6d361a2d43d7590da657541b682947a0/pipeline_waifu.py:13\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# Transformer\u001b[39;00m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdiffusers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m SanaTransformer2DModel\n\u001b[0;32m---> 13\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mWaifuPipeline\u001b[39;00m(\u001b[43mDiffusionPipeline\u001b[49m):\n\u001b[1;32m     14\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     15\u001b[0m \u001b[38;5;124;03m    Pipeline for text-to-image generation using [waifu](https://github.com/recoilme/waifu).\u001b[39;00m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m     18\u001b[0m     model_cpu_offload_seq \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtext_encoder->transformer->vae\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
-      "\u001b[0;31mNameError\u001b[0m: name 'DiffusionPipeline' is not defined"
+      "File \u001b[0;32m~/.cache/huggingface/modules/diffusers_modules/local/AiArtLab--waifu-2b/198e5ef5b8669576c07d632d8623f6da0f35d61e/pipeline_waifu.py:17\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# Transformer\u001b[39;00m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdiffusers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m SanaTransformer2DModel\n\u001b[0;32m---> 17\u001b[0m \u001b[38;5;28;43;01mclass\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43;01mWaifuPipeline\u001b[39;49;00m\u001b[43m(\u001b[49m\u001b[43mDiffusionPipeline\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;250;43m    \u001b[39;49m\u001b[38;5;124;43mr\u001b[39;49m\u001b[38;5;124;43;03m\"\"\"\u001b[39;49;00m\n\u001b[1;32m     19\u001b[0m \u001b[38;5;124;43;03m    Pipeline for text-to-image generation using [waifu](https://github.com/recoilme/waifu).\u001b[39;49;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;124;43;03m    \"\"\"\u001b[39;49;00m\n\u001b[1;32m     22\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel_cpu_offload_seq\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mtext_encoder->transformer->vae\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\n",
+      "File \u001b[0;32m~/.cache/huggingface/modules/diffusers_modules/local/AiArtLab--waifu-2b/198e5ef5b8669576c07d632d8623f6da0f35d61e/pipeline_waifu.py:45\u001b[0m, in \u001b[0;36mWaifuPipeline\u001b[0;34m()\u001b[0m\n\u001b[1;32m     39\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvae_scale_factor \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m8\u001b[39m\n\u001b[1;32m     40\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mimage_processor \u001b[38;5;241m=\u001b[39m PixArtImageProcessor(vae_scale_factor\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvae_scale_factor)\n\u001b[1;32m     42\u001b[0m \u001b[38;5;129m@torch\u001b[39m\u001b[38;5;241m.\u001b[39mno_grad()\n\u001b[1;32m     43\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\n\u001b[1;32m     44\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m---> 45\u001b[0m     prompt: \u001b[43mUnion\u001b[49m[\u001b[38;5;28mstr\u001b[39m, List[\u001b[38;5;28mstr\u001b[39m]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     46\u001b[0m     negative_prompt: \u001b[38;5;28mstr\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     47\u001b[0m     num_inference_steps: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m20\u001b[39m,\n\u001b[1;32m     48\u001b[0m     timesteps: List[\u001b[38;5;28mint\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     49\u001b[0m     sigmas: List[\u001b[38;5;28mfloat\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     50\u001b[0m     guidance_scale: \u001b[38;5;28mfloat\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m4.5\u001b[39m,\n\u001b[1;32m     51\u001b[0m     num_images_per_prompt: Optional[\u001b[38;5;28mint\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     52\u001b[0m     height: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m512\u001b[39m,\n\u001b[1;32m     53\u001b[0m     width: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m512\u001b[39m,\n\u001b[1;32m     54\u001b[0m     eta: \u001b[38;5;28mfloat\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.0\u001b[39m,\n\u001b[1;32m     55\u001b[0m     generator: Optional[Union[torch\u001b[38;5;241m.\u001b[39mGenerator, List[torch\u001b[38;5;241m.\u001b[39mGenerator]]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     56\u001b[0m     latents: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     57\u001b[0m     prompt_embeds: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     58\u001b[0m     prompt_attention_mask: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     59\u001b[0m     negative_prompt_embeds: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     60\u001b[0m     negative_prompt_attention_mask: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     61\u001b[0m     output_type: Optional[\u001b[38;5;28mstr\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpil\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     62\u001b[0m     return_dict: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     63\u001b[0m     callback_on_step_end: Optional[Callable[[\u001b[38;5;28mint\u001b[39m, \u001b[38;5;28mint\u001b[39m, Dict], \u001b[38;5;28;01mNone\u001b[39;00m]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     64\u001b[0m     callback_on_step_end_tensor_inputs: List[\u001b[38;5;28mstr\u001b[39m] \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mlatents\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m     65\u001b[0m     max_sequence_length: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m512\u001b[39m,\n\u001b[1;32m     66\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[SanaPipelineOutput, Tuple]:\n\u001b[1;32m     67\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     68\u001b[0m \u001b[38;5;124;03m    Function invoked when calling the pipeline for generation.\u001b[39;00m\n\u001b[1;32m     69\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;124;03m            otherwise a `tuple` is returned where the first element is a list with the generated images\u001b[39;00m\n\u001b[1;32m    143\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m    145\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'Union' is not defined"
      ]
     }
    ],

pipeline_waifu.py

@@ -1,5 +1,6 @@
 import torch
 from diffusers import DiffusionPipeline
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 # waifu
 # tokenizer