You can deploy this model and start earning money today!

XXXXingyu/
dreamir

Name: dreamir
Availability: InStock

Views

No views yet

DreamIR · Image Restoration (IR) — Inference

0. 硬件与环境要求

单卡 NVIDIA GPU（推荐 A100/H100 40G+ 显存；跑 512×512 输入约需 ~30GB 显存，因为整个 DiT 常驻 GPU 且未启用 VRAM offload）
CUDA 11.8（对齐下方 torch wheel）；如需其他 CUDA 版本，请自行替换对应的 torch 安装命令
Linux + Python 3.10
磁盘：base model 最小下载 ≈ 12 MB（仅 processor/tokenizer 相关文件）或完整下载 ≈ 35 GB；IR checkpoint ≈ 14 GB（bf16 dit 权重）

1. 安装环境

bash

conda create -n diffsynth python=3.10 -y
conda activate diffsynth

# 请根据你机器上的 CUDA 版本调整下面这条命令
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt

requirements.txt 包含的关键依赖：safetensors, transformers, einops, Pillow, numpy, tqdm, modelscope。

2. 下载所需权重

推理需要两份权重：

#	名称	用途	来源	大小
A	HiDream-O1-Image（base model）	提供 tokenizer + image processor（`AutoProcessor.from_pretrained`）	ModelScope / HuggingFace	~35 GB（含权重）；若只要 tokenizer/processor 约 12 MB
B	IR fine-tuned checkpoint	DiT 网络的 IR 微调权重，直接覆盖 DiT 全部参数	内部产出（本仓库 fine-tune 得到的 `step-*.safetensors`）	~14 GB

2.1 下载 A：HiDream-O1-Image（base model 目录）

代码里 AutoProcessor.from_pretrained(base_model_dir) 只用到 tokenizer + image processor 的配置， 并不加载 shard 权重。你有两种选择：

方式 A1：最小下载（推荐，省 ~35 GB 磁盘）

只下载 processor / tokenizer 相关的小文件（推理只会用到它们）：

bash

export BASE_MODEL_DIR=/path/to/hidream_o1_image

modelscope download --model HiDream-ai/HiDream-O1-Image --local_dir $BASE_MODEL_DIR \
  config.json generation_config.json chat_template.json \
  preprocessor_config.json video_preprocessor_config.json \
  tokenizer_config.json tokenizer.json vocab.json merges.txt

说明：run_inference.py 不会加载 base model 的 DiT/VAE 权重，DiT 权重完全由步骤 B 的 IR checkpoint 提供。因此这种最小下载方式（约 12 MB）足以支撑本仓库的 IR 推理。

方式 A2：完整下载（简单省事，但会占 ~35 GB 磁盘）

用 ModelScope（国内速度快，且本仓库 requirements 已含 modelscope）：

bash

# 目标目录，随意起名，后面 run_inference.py 里要指向它
export BASE_MODEL_DIR=/path/to/hidream_o1_image

# 通过 modelscope CLI 下载
modelscope download --model HiDream-ai/HiDream-O1-Image --local_dir $BASE_MODEL_DIR

或者用 huggingface（需要能访问 huggingface.co）：

bash

pip install -U "huggingface_hub[cli]"
huggingface-cli download HiDream-ai/HiDream-O1-Image --local-dir $BASE_MODEL_DIR

下载完成后目录里应包含如下文件（示例，实际以仓库为准）：

$BASE_MODEL_DIR/
├── config.json
├── generation_config.json
├── chat_template.json
├── preprocessor_config.json
├── video_preprocessor_config.json
├── tokenizer_config.json
├── tokenizer.json
├── vocab.json
├── merges.txt
├── model.safetensors.index.json
├── model-00001-of-00008.safetensors  ← 这些 shards 本次 IR 推理其实用不上
├── model-00002-of-00008.safetensors
├── ...
└── model-00008-of-00008.safetensors

2.2 下载 B：IR checkpoint

把权重放到你选择的路径，

export IR_CKPT=/path/to/dreamIR/dreamir.safetensors

3. 配置路径

打开 run_inference.py，把顶部这两行改成你自己的路径：

python

# run_inference.py, line 21-22
base_model_dir = "/path/to/hidream_o1_image"                  # ← 步骤 2.1 下载到的目录
ir_checkpoint  = "/path/to/dreamIR/dreamir.safetensors"  # ← 步骤 2.2 拿到的文件

其他可按需修改：

output_folder（默认 results/）
DEFAULT_INPUT_FOLDER（也可以命令行 --input_folder 传入，见下文）
PROMPT（IR restoration 用的固定 prompt，一般不需要改）

4. 运行 Inference

准备一个装了 LQ 输入图的文件夹（支持 .jpg / .jpeg / .png，脚本目前的实现假设图像为 512×512，其他尺寸会按图像自身尺寸跑，但需保证宽高能被 patch size 整除）：

bash

YOUR_LQ_DATA_FOLDER_PATH=/path/to/your/LQ_images

python run_inference.py --input_folder $YOUR_LQ_DATA_FOLDER_PATH

首次运行会看到类似如下日志：

Instantiating DiT architecture directly on cuda in torch.bfloat16...
Loading IR checkpoint directly to cuda: /path/to/dreamir.safetensors
Loading processor (tokenizer + image preprocessor)...
Pipeline ready.
Found N images. Saving to: results
IR Processing: 100%|██████████| N/N [xx:xx<00:00, ...]
Done. Saved to: results

修复后的图像会以 JPEG（quality=96, 4:4:4）保存到 output_folder 中，同名文件（后缀改 .jpg）。

5. 目录结构参考

dreamir/
├── README.md                 # 本文件
├── requirements.txt          # Python 依赖
├── run_inference.py          # 推理入口脚本
├── diffsynth/                # 精简过的 DiffSynth-Studio 代码（HiDream-O1-Image 相关模块）
│   ├── pipelines/hidream_o1_image.py
│   ├── models/
│   │   ├── hidream_o1_image_dit.py
│   │   └── hidream_common.py
│   └── ...
└── results/                  # 推理产物默认保存目录

6. 速跑清单

bash

# 0. 装环境
conda create -n diffsynth python=3.10 -y && conda activate diffsynth
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

# 1. 下载 base model 里 processor/tokenizer 相关的小文件（推理只需这些，约 12 MB）
export BASE_MODEL_DIR=$PWD/checkpoints/hidream_o1_image
modelscope download --model HiDream-ai/HiDream-O1-Image --local_dir $BASE_MODEL_DIR \
  config.json generation_config.json chat_template.json \
  preprocessor_config.json video_preprocessor_config.json \
  tokenizer_config.json tokenizer.json vocab.json merges.txt

# 2. 拿到 IR checkpoint（内部/自训练）
export IR_CKPT=/path/to/dreamir.safetensors

# 3. 把上面两个变量的值填进 run_inference.py 的 base_model_dir / ir_checkpoint

# 4. 跑！
python run_inference.py --input_folder /path/to/your/LQ_images

XXXXingyu/dreamir

DreamIR · Image Restoration (IR) — Inference

0. 硬件与环境要求

1. 安装环境

2. 下载所需权重

2.1 下载 A：HiDream-O1-Image（base model 目录）

方式 A1：最小下载（推荐，省 ~35 GB 磁盘）

方式 A2：完整下载（简单省事，但会占 ~35 GB 磁盘）

2.2 下载 B：IR checkpoint

3. 配置路径

4. 运行 Inference

5. 目录结构参考

6. 速跑清单

XXXXingyu/
dreamir