tt_hidream_o1 — Python reference
HiDream-O1-Image (undistilled) on one Tenstorrent Blackhole p100a. The Qwen3-VL-8B language model runs on the card; the vision tower, the sampler and the patch maths are host. There is no VAE — the model works directly on 32×32 pixel patches.
Install
pip install -e . # inside a tt-metal / ttnn environment; ttnn is not on PyPI
pip install -e ".[server]" # also fastapi / uvicorn / pydantic
Weights
HiDream-ai/HiDream-O1-Image, pinned to 0b0901d99f200389e138c61946af1185f5f49a13, downloaded into
the HF cache on first use. HIDREAM_SNAPSHOT points at a directory you manage instead.
API
open_device(trace_region_size=0, l1_small_size=32768)
Opens a 1×1 mesh with DispatchCoreType.WORKER. Pair with close_device(dev).
HiDreamTT(dev, snapshot=None, prec=None)
Loads the language model onto the card. The vision tower is built lazily on first use.
HiDreamTT.generate(prompt, width, height, seed, refs=None, steps=50, guidance_scale=5.0, snap=True, keep_original_aspect=False) -> (PIL.Image, Request, Timing)
| Argument | Meaning |
|---|---|
width, height |
On a 32 px grid. Other sizes are generated larger and center-cropped. |
snap |
Pass snap=False for the size you asked for. The upstream default rounds the request to the official 4–5 MP resolution list, so 1024² becomes 2048². |
refs |
Reference images for the editing path. Only usable at the official ~4 MP sizes — see the caveat below. |
steps, guidance_scale |
50 and 5.0. This is the undistilled model: 50 steps at CFG 5 is 100 forwards. |
keep_original_aspect |
Sizes the target from the reference, scaled to 2048² area. |
The middle value is the resolved Request: the size, patch grid and step count the pipeline
actually used after snap and the aspect rules were applied. Most callers want the first and
third and bind it to _req; the parity checks in experiments/ read it to rebuild the same
request on the CPU.
Timing.values holds vision_s, prefix_s, lm_s.
Lower-level entry points
| Method | |
|---|---|
branches(req, timing) |
builds the conditional and unconditional prefixes |
denoise(req, seed, steps, guidance_scale, timing) |
the sampler loop |
dram_stats(dev) -> dict
total_bytes, allocated_bytes, free_bytes.
The editing caveat
refs works, but only at the official ~4 MP sizes. At 512–640 px every edit comes out as noise — and
so does the official CPU pipeline at that size, so this is upstream behaviour, not a porting bug
(experiments/diag_edit.py). At 2048² the edit is correct and matches the CPU reference to PCC 0.9865,
but one costs about 6 minutes on the card.
Modules
| Module | Role |
|---|---|
pipeline.py |
HiDreamTT, device open/close, timing |
host.py |
host maths: the chat template, mRoPE positions, patching, the schedule, request construction |
tt_lm.py |
Qwen3-VL-8B language model (36 layers) on the card |
ref_algo.py |
torch reference implementations used by the verification scripts |
upstream/ |
official model code, vendored at 2c2d29f (MIT), with flash-attn made optional and a CUDA call guarded |
config.py |
static shapes, the pinned revision, snapshot resolution |
server.py |
the HTTP app |