Birch-san

Installing & running cene's imagen-pytorch

Install Anaconda via command-line arm64 installer

Install PyTorch nightly (to get M1 GPU support)

conda update -n base -c defaults conda
conda create -n torch-nightly python=3.9
pip install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Birch-san

import torch as th

from itertools import tee

from PIL import Image
from IPython.display import display

from imagen_pytorch.model_creation import create_model_and_diffusion as create_model_and_diffusion_imagen
from imagen_pytorch.model_creation import model_and_diffusion_defaults as model_and_diffusion_defaults_imagen
from imagen_pytorch.train_all import _fix_path

Birch-san

# https://google.github.io/iree/building-from-source/getting-started/#reference-pages
# https://github.com/google/iree/blob/main/docs/developers/get_started/cmake_options_and_variables.md
git clone --branch scatter --depth 1 https://github.com/Birch-san/iree.git
git submodule update --init
cd iree
pip install pybind11
export CMAKE_OSX_ARCHITECTURES="arm64"
export MACOSX_DEPLOYMENT_TARGET=11.0
export IREE_ENABLE_CPUINFO=OFF
export IREE_HAL_DRIVER_VULKAN=ON

Birch-san

# see https://github.com/google/iree/issues/9361,
#     https://github.com/google/iree/pull/9378

from iree import compiler, runtime as rt
rt.flags.parse_flags("--task_topology_group_count=8")

CODE = """
#loc0 = loc(unknown)
module @jit_prim_fun.12 {
  func.func public @main(%arg0: tensor<1x1xi32> loc(unknown), %arg1: tensor<1xi32> loc(unknown), %arg2: tensor<1xi32> loc(unknown)) -> tensor<1x1xi32> {

Birch-san

from iree import compiler

CODE = """
#loc0 = loc(unknown)
module @jit__threefry_split.11 {
  func.func public @main(%arg0: tensor<2xui32> loc(unknown)) -> tensor<2x2xui32> {
    %0 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32> loc(#loc1)
    %1 = "mhlo.slice"(%arg0) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32> loc(#loc2)
    %2 = "mhlo.reshape"(%1) : (tensor<1xui32>) -> tensor<ui32> loc(#loc3)
    %3 = "mhlo.slice"(%arg0) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32> loc(#loc4)

Birch-san

#loc0 = loc(unknown)
module @pmap_p_generate.14 {
  func.func public @main(%arg0: tensor<1x64xi32> loc(unknown), %arg1: tensor<1x64xi32> loc(unknown), %arg2: tensor<1x64xi32> loc(unknown), %arg3: tensor<1x64xi32> loc(unknown), %arg4: tensor<2xui32> loc(unknown), %arg5: tensor<2048x16416xf32> loc(unknown), %arg6: tensor<256x2048xf32> loc(unknown), %arg7: tensor<16416x2048xf32> loc(unknown), %arg8: tensor<2048xf32> loc(unknown), %arg9: tensor<2048xf32> loc(unknown), %arg10: tensor<2048xf32> loc(unknown), %arg11: tensor<24x2048x2048xf32> loc(unknown), %arg12: tensor<24x2048x2048xf32> loc(unknown), %arg13: tensor<24x2048x2048xf32> loc(unknown), %arg14: tensor<24x2048x2048xf32> loc(unknown), %arg15: tensor<24x2048x2048xf32> loc(unknown), %arg16: tensor<24x2048x2048xf32> loc(unknown), %arg17: tensor<24x2048x2048xf32> loc(unknown), %arg18: tensor<24x2048x2048xf32> loc(unknown), %arg19: tensor<24x2048x4096xf32> loc(unknown), %arg20: tensor<24x2048x4096xf32> loc(unknown), %arg21: tensor<24x4096x2048xf32> loc(unknown),

Birch-san

A_chan.gif[FRAMES:]6[GENERAL_PICKN:]smile		blue_bow	bow	hair_bow	black_jacket	id_card	lanyard	solo	glasses	purple_hair	bangs	black-framed_eyewear	1girl	official_alternate_costume	short_hair	transparent_background	closed_mouth	chibi[GENERAL_CRUCIAL:]typing	laptop	upper_body[PROPER_NOUNS:]a-chan_(hololive)	walfie_(style)
a_way_out.gif[FRAMES:]5[GENERAL_PICKN:]brown_footwear	blonde_hair	skirt	brown_skirt	closed_mouth	necktie	red_necktie	shirt	white_shirt	smile	short_hair	transparent_background	kneehighs	shark_tail	hood	hood_down	hoodie	shark_girl	silver_hair	tail	fish_tail	chibi	blue_hoodie		grey_hair	long_hair	grey_legwear	long_sleeves	cabbie_hat	bangs	plaid_skirt[GENERAL_CRUCIAL:]2girls	multiple_girls	locked_arms	abseiling	escape	falling	full_body	|_|	._.[PROPER_NOUNS:]smol_gura	gawr_gura	smol_ame	watson_amelia	walfie_(style)
Ame_Birthday_2022.gif[FRAMES:]30[GENERAL_PICKN:]brown_footwear	blonde_hair	skirt	brown_skirt	closed_mouth	necktie	red_necktie	shirt	white_shirt	smile	short_hair	kneehighs	grey_legwear	tra

Birch-san

::a-chan_(hololive)
blue_bow	bow	hair_bow	short_hair	black-framed_eyewear
black_jacket	id_card	lanyard
official_alternate_costume

::smol_ame	watson_amelia
brown_footwear	blonde_hair	skirt	brown_skirt	necktie	red_necktie	shirt	white_shirt	short_hair	kneehighs	grey_legwear	long_sleeves	cabbie_hat	bangs	plaid_skirt
brown_jacket

::smol_gura	gawr_gura

Birch-san

from torch.nn import Embedding

from typing import Tuple, TypeVar, Iterable
from typing_extensions import TypeAlias
from enum import Enum, auto
from math import ceil
from torch import BoolTensor, LongTensor, sparse_coo_tensor, ones
from itertools import chain

class Label(Enum):

Birch-san

supported sigmas:
tensor([ 0.0292,  0.0413,  0.0507,  0.0586,  0.0656,  0.0720,  0.0779,  0.0834,
         0.0886,  0.0936,  0.0983,  0.1028,  0.1072,  0.1114,  0.1155,  0.1195,
         0.1234,  0.1271,  0.1308,  0.1345,  0.1380,  0.1415,  0.1449,  0.1482,
         0.1515,  0.1548,  0.1580,  0.1611,  0.1642,  0.1673,  0.1704,  0.1734,
         0.1763,  0.1793,  0.1822,  0.1850,  0.1879,  0.1907,  0.1935,  0.1963,
         0.1991,  0.2018,  0.2045,  0.2072,  0.2099,  0.2125,  0.2152,  0.2178,
         0.2204,  0.2230,  0.2256,  0.2281,  0.2307,  0.2332,  0.2358,  0.2383,
         0.2408,  0.2433,  0.2458,  0.2482,  0.2507,  0.2531,  0.2556,  0.2580,
         0.2604,  0.2628,  0.2653,  0.2677,  0.2700,  0.2724,  0.2748,  0.2772,