HiDream-O1: support area conditioning

[kijai]

Position IDs and vinput_mask were precomputed in extra_conds against the full noise shape, but area conds crop the latent before forward, the precomputed layout then mismatched the cropped x and errored out.

Fix: pass the area-cropped shape into build_extra_conds upfront so the sequence layout is built correctly.

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request adds area-based noise cropping to the HiDreamO1.extra_conds() method. When an area argument is provided, the method computes cropped height and width dimensions (clamped to the existing noise spatial bounds), allocates a new tensor containing only the cropped region with dimensions [3, crop_h, crop_w] while preserving dtype and device, and uses this cropped tensor for subsequent conditioning construction.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'HiDream-O1: support area conditioning' clearly and directly summarizes the main change: adding area conditioning support to the HiDream-O1 model.
Description check	✅ Passed	The description explains the specific problem (position IDs and vinput_mask precomputed against full shape causing mismatches) and the solution (passing area-cropped shape into build_extra_conds), which aligns with the changeset.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

comfy/model_base.py (1)

1694-1700: ⚖️ Poor tradeoff

Consider using torch.zeros() instead of torch.empty() for clarity and safety.

On line 1699, torch.empty() creates an uninitialized tensor. While build_extra_conds() only reads the shape, device, and dtype properties (never the actual tensor values), using torch.zeros() would be more conventional and self-documenting—making it explicit that the tensor contents are not used. This also protects against future maintenance issues if the code is modified.

-    noise = torch.empty((noise.shape[0], 3, crop_h, crop_w), dtype=noise.dtype, device=noise.device)
+    noise = torch.zeros((noise.shape[0], 3, crop_h, crop_w), dtype=noise.dtype, device=noise.device)

The performance impact is negligible for small tensors used only for shape references.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@comfy/model_base.py` around lines 1694 - 1700, The code creates an
uninitialized tensor with torch.empty(...) for the area-handling branch (when
area is not None) using the local variables noise, crop_h, crop_w; change that
torch.empty(...) call to torch.zeros(...) while preserving the same shape,
dtype, and device so the new tensor is zero-initialized (use the existing
noise.shape[0], 3, crop_h, crop_w, noise.dtype, noise.device), updating the
occurrence in the area-handling block (the lines referencing area, crop_h,
crop_w and the torch.empty call inside build_extra_conds / model_base.py).

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@comfy/model_base.py`:
- Around line 1694-1700: The code creates an uninitialized tensor with
torch.empty(...) for the area-handling branch (when area is not None) using the
local variables noise, crop_h, crop_w; change that torch.empty(...) call to
torch.zeros(...) while preserving the same shape, dtype, and device so the new
tensor is zero-initialized (use the existing noise.shape[0], 3, crop_h, crop_w,
noise.dtype, noise.device), updating the occurrence in the area-handling block
(the lines referencing area, crop_h, crop_w and the torch.empty call inside
build_extra_conds / model_base.py).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a8aa4469-83ab-42ac-8629-dc51d9b0cea8

📥 Commits

Reviewing files that changed from the base of the PR and between 7c4d95d and 4682039.

📒 Files selected for processing (1)

• comfy/model_base.py