| # 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 |
|
|
| ```bash |
| 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 | |
|
|
