画像生成¶

tensai image は Qwen-Image-2.1 でプロンプトから絵を描きます。ほかと同じ純粋な Go で、cgo なし、C コンパイラ不要、XOR ネットワークを学習させるのと同じカーネルの上で動きます。

$ tensai image -size 256 -steps 20 "a calico cat asleep on a stack of books"
prompt: 18 tokens in 6s
transformer: loaded in 0s
step 20/20 in 3m39s
wrote out.png, 256x256

何が走るか¶

3 つのモデルが順に走ります。どれも 7GB あるので、同時に持つのは 1 つだけです。

プロンプトエンコーダ。denoising transformer はプロンプトをトークンではなく vision-language モデルの隠れ状態として読むので、まず Qwen3-VL 8B の言語側 36 層が走ります。渡すのは最終ノルムを通す前の最後の層の出力で、テンプレートのシステム前置きは落とします。終わったら解放します。
denoising transformer。プロンプトのトークンと画像の latent 格子をつないだ 1 本の列を、単流のブロック 32 層が通ります。1 ブロックは attention と SwiGLU の feed-forward で、挟む層正規化は重みを持ちません。持たないスケールとゲートは timestep から来ます。これを 20 回通すと latent がノイズから絵になります。
デコーダ。latent の 1 点を 16 ピクセルの正方形にする 2D オートエンコーダです。チェックポイントがそのまま生成するアルファチャンネルも持ちます。

すべてのピースを参照実装と突き合わせてあります (transformer・スケジュール・デコーダは diffusers、プロンプトエンコーダは transformers)。差は 6.5e-8 から 2.4e-4 で、float32 の丸め誤差の範囲です。参照は全層ではメモリに載らないため層数を絞って照合していますが、ブロック以外の配線は層数に依存しません。

チェックポイントの入手¶

tensai image "a calico cat asleep on a stack of books"
tensai image -model Comfy-Org/Qwen-Image-2.1 "a calico cat asleep on a stack of books"

-model は run や chat と同じ形でチェックポイントを指定します。リポジトリ名、tensai models が表示する名前、パスのどれでも受け付けます。次の 2 つのリポジトリは初回に自動でダウンロードします。中断などでファイルが欠けていれば、次の実行でそのファイルだけを取り直します。すべて揃っていれば通信はしません。

Qwen/Qwen-Image-2.1 (既定): diffusers 形式のチェックポイントで、およそ 31GB です。~/.cache/tensai/Qwen/Qwen-Image-2.1 に置きます (組織名を付ける前のバージョンで取得した ~/.cache/tensai/Qwen-Image-2.1 があればそれを使います)。重みファイル名は各コンポーネントの index から読むので、リポジトリ側で分割数が変わっても解決できます。
Comfy-Org/Qwen-Image-2.1: ComfyUI 向けに詰め直したもので、~/.cache/tensai/Comfy-Org/Qwen-Image-2.1 に置きます。transformer とプロンプトエンコーダが最初から 8 ビットなので、int8 版の合計はおよそ 17GB です。重みは入力 256 列ごとに Hadamard 行列で回転させてから、行ごとのスケール付き int8 にしてあります。ローダーはこの 2 つを戻してから自前の量子化をかけるので、どちらのチェックポイントも同じコードで動きます。ComfyUI 版にはトークナイザと VAE の潜在統計が含まれないため、この小さな 2 ファイルだけは Qwen のリポジトリから取ります。

パスにはどちらの配置のディレクトリも指定できます。diffusers なら text_encoder・transformer・vae・processor、ComfyUI なら text_encoders・diffusion_models・vae に processor/tokenizer.json と vae/config.json を加えたものです。ComfyUI の配置では _int8_convrot のファイルを選び、無ければ _bf16 を使います。以前ダウンロードに使っていた -fetch は、指定しても何もしません。

幅とメモリ¶

チェックポイントは bfloat16 で、transformer だけで 14GB あります。そのため読み込みながら量子化します。既定は 8 ビット、-q4 で 4 ビットです。

幅	transformer	velocity 誤差 (8 ブロック)
`-f32`	28 GB	—
既定	7.0 GB	9.1%
`-q4`	3.5 GB	13.5%

誤差は同じモデルの float 版との比較で、深さに対して掛け算では増えません。1 ブロックで 8 ビットが 7.2%、4 ビットが 12.8%、8 ブロックで 9.1% と 13.5% です。どちらもプロンプトどおりの絵になり、4 ビットは細部の解釈が変わります。

上の数字は回転を入れる前のものです。今は量子化の前に、すべての重みを入力 256 列ごとに Hadamard 行列で回転させ、積の入力にも同じ回転を掛けます。行列は直交なので積の結果は変わりませんが、行のスケールを決めてしまう少数の外れ値の列がグループ全体にならされます。8 ビットが一番損をしていたのはそこです。ComfyUI の int8 ファイルと同じ方式なので、そちらの重みは最初からこの形で届きます。参照入力での最初のブロックでは、8 ビットの誤差が 6.95% から 3.48% とほぼ半分になり、4 ビットは 19.3% から 16.2% になります。モデル全体の 4 ブロック分で見ると、8 ビットの velocity 誤差は CPU で 7.8% から 2.4%、-gpu で 4.8% から 2.1% に下がります。デバイス上では feed-forward の出力側の射影が読む値もデバイスで作られるので、回転もデバイスで掛けます。行を 256 ずつのグループとして見て H との積を 1 回計算するだけで、射影そのものの数パーセントのコストです。

70 億パラメータの量子化には数分かかり、結果は毎回同じなので、チェックポイントの隣に tensai-q8.cache (または tensai-q4.cache) として書き、次回は mmap で読みます。初回は遅く、2 回目からは速くなります。プロンプトエンコーダの起動が 5 分 23 秒から 7 秒、transformer が 7 分 7 秒から 2 秒です。キャッシュは隣の .safetensors の名前・サイズ・更新時刻で紐付けてあるので、再ダウンロードすれば作り直します。

サイズと時間¶

-size はピクセル単位で、32 の倍数に切り下げます (デコーダが 16 ピクセル単位で、チェックポイントの配置がそれを 2 つずつ束ねるため)。16 コアのノート PC、AVX2 ビルドでの実測です。

サイズ	latent トークン	1 ステップ	20 ステップ
256x256	256	11 秒	3 分 46 秒
512x512	1024	50 秒	約 17 分

トークン数の増加より 1 ステップの伸びが大きいのは、attention が列長の二乗だからです。

GPU¶

-tags wgpu24 付きのビルドで -gpu を指定すると、各ブロックの feed-forward がデバイスに載ります。3 つの射影とゲートで、1 ブロックの計算の 7 割を占め、位置スキームもマスクも要求しない部分です。512x512 の 1 ステップが 50 秒から 31 秒になり、しかも float の重みが言う答えに近づきます (CPU の 7.8% に対して 4.7%)。デバイス側のほうが各積の活性化を細かく量子化するためです。

Attention のスコア計算・softmax・値の集約も GPU で実行します。テキスト部分は因果マスク、画像部分は全トークン参照として処理し、スコア用のメモリを抑えるためクエリを分割します。 K/V はブロックごとに一度だけ転送し、分割したクエリ間で共有します。 Attention の射影と各種ノルム、3 軸の rotary embedding は CPU に残ります。FFN の GPU 演算はまとめて送信します。追加の常駐重みは必要ありません。

重みは実行中ずっとデバイスに置きます。気を使うのはそこです。このデバイスは 1 バッファ 128MiB が上限で (int8 の重みは収まり float は収まりません)、常駐がおよそ 5.5GB を超えるとドライバがデバイスを落とします。しかも静かに落ちます。確保は成功を返し続け、後からプロセスが倒れます。それを知らせる仕組みが無いので、-gpu-budget がアップロード前に合計を数えて、そこで止めます。8 ビットの feed-forward は 4.5GiB で既定の 4 を超えるため、-gpu 単独では予算に入るブロックだけを載せ、残りは CPU に置いたまま、何ブロック載ったかを出します。全部載せるには -q4 (2.3GiB) か予算の引き上げです。

フラグ¶

tensai image [flags] <prompt>

  -model string   リポジトリ名、キャッシュ内の名前、またはパス (既定 "Qwen/Qwen-Image-2.1")
  -o string       書き出し先 (既定 "out.png")
  -size int       幅と高さ、ピクセル (既定 256)
  -steps int      デノイジングのステップ数 (既定 20)
  -seed int       ノイズのシード (既定 1)
  -q4             8 ビットではなく 4 ビットに量子化する
  -f32            重みを float のまま持つ。およそ 42GB 必要
  -negative str   避けたいもの。-cfg を 1 より大きくする必要がある
  -cfg float      どれだけ避けるか (既定 1、off)
  -gpu            feed-forward と attention を GPU で走らせる
  -gpu-budget num GPU に載せてよい重みの GB 数 (既定 4)
  -gpu-projections attention の射影を GPU にストリームする (-gpu で有効)
  -gpu-vae        デコーダの畳み込みを GPU で走らせる (-gpu で有効)
  -cpuprofile str CPU プロファイルの保存先
  -q              エラー以外を出さない

実行ログには GPU への重み読み込み、transformer の解放、デコーダの読み込みとデコードの時間も表示します。step の時間は生成開始からの累計です。

ガイダンス¶

-cfg を 1 より大きくすると classifier-free guidance が有効になります。各ステップで transformer に 2 回聞き (プロンプトが求めるものと -negative が求めるもの)、その差を前者の側に伸ばします。色が分離し質感が出て絵が締まりますが、1 ステップのコストが 2 倍になるため既定では off です。

tensai image -cfg 4 -negative "blurry, low quality, watermark" \
  "a calico cat asleep on a stack of books"

256x256 の 20 ステップが 3 分 46 秒から 6 分 29 秒になります。2 つのプロンプトはエンコーダの 1 回の読み込みで両方処理するので、そちらの追加コストはありません。

まだ無いもの¶

text-to-image だけです。チェックポイントは画像編集もできますが、そちらは未実装です。プロンプトエンコーダは言語側だけを走らせるので vision tower は使いません。参照実装がステップ間で使い回す key-value キャッシュ (プロンプトは timestep 0 から変調するので全ステップで変わりません) も作っていません。