helios model/pipeline review

# `helios` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Duplicate search status: searched GitHub Issues/PRs for `helios`, `HeliosPipeline latents`, `Helios num_videos_per_prompt`, `helios ftfy`, `helios timestep`, `prepare_video_latents first_frame_latent`, and `HeliosPyramidPipeline test`. No specific duplicate issue/PR was found. PR https://github.com/huggingface/diffusers/pull/13218 only skips a Helios float16 save/load test and is not a duplicate.

## Issue 1: `num_videos_per_prompt` breaks Helios generation

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L599-L613
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L703-L711
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L792-L802
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/before_denoise.py#L100-L116
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/before_denoise.py#L666-L688

Problem:
`encode_prompt` expands prompt embeddings to `batch_size * num_videos_per_prompt`, but latent/history tensors are still allocated with the original prompt batch size. The first transformer call then receives hidden states with batch size 1 and prompt embeddings with batch size 2.

Impact:
Users cannot request multiple videos per prompt. The modular pipeline has the same contract problem: `HeliosTextInputStep` expands embeddings but leaves `batch_size` as the pre-expansion value used by history and latent prep.

Reproduction:
```python
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, HeliosPipeline, HeliosScheduler, HeliosTransformer3DModel

vae = AutoencoderKLWan(base_dim=3, z_dim=16, dim_mult=[1, 1, 1, 1], num_res_blocks=1, temperal_downsample=[False, True, True])
scheduler = HeliosScheduler(stage_range=[0, 1], stages=1, use_dynamic_shifting=True)
config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
pipe = HeliosPipeline(
transformer=HeliosTransformer3DModel(
patch_size=(1, 2, 2), num_attention_heads=2, attention_head_dim=12, in_channels=16, out_channels=16,
text_dim=32, freq_dim=256, ffn_dim=32, num_layers=2, rope_dim=(4, 4, 4),
),
vae=vae,
scheduler=scheduler,
text_encoder=T5EncoderModel(config),
tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"),
).to("cpu")
pipe.set_progress_bar_config(disable=True)

pipe(
prompt="dance monkey",
negative_prompt="negative",
generator=torch.Generator("cpu").manual_seed(0),
num_inference_steps=1,
guidance_scale=1.0,
height=16,
width=16,
num_frames=9,
max_sequence_length=16,
output_type="latent",
num_videos_per_prompt=2,
)
```

Relevant precedent:
`WanPipeline` and other video pipelines propagate the effective batch size into latent prep after prompt/image expansion.

Suggested fix:
```python
effective_batch_size = batch_size * num_videos_per_prompt

history_latents = torch.zeros(
effective_batch_size,
num_channels_latents,
num_history_latent_frames,
height // self.vae_scale_factor_spatial,
width // self.vae_scale_factor_spatial,
device=device,
dtype=torch.float32,
)

latents = self.prepare_latents(
effective_batch_size,
num_channels_latents,
height,
width,
window_num_frames,
dtype=torch.float32,
device=device,
generator=generator,
latents=chunk_latents,
)
```

For modular Helios, set `block_state.batch_size` to the expanded effective batch size or add a separate `effective_batch_size` output and use it for history/latent allocation.

## Issue 2: Public `latents` input is ignored

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L792-L802
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py#L844-L854
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/denoise.py#L243-L265
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/denoise.py#L268-L313

Problem:
Both standard pipelines accept `latents`, document it as pre-generated noisy latents, and list it in required optional params, but the denoise loop always calls `prepare_latents(..., latents=None)`. The modular noise blocks also always sample new noise.

Impact:
Users cannot reproduce or edit a generation by supplying their own initial noise. Tests can miss this because the signature exists and generation still succeeds.

Reproduction:
```python
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import AutoencoderKLWan, HeliosPipeline, HeliosScheduler, HeliosTransformer3DModel

def make_pipe():
vae = AutoencoderKLWan(base_dim=3, z_dim=16, dim_mult=[1, 1, 1, 1], num_res_blocks=1, temperal_downsample=[False, True, True])
config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
return HeliosPipeline(
transformer=HeliosTransformer3DModel(
patch_size=(1, 2, 2), num_attention_heads=2, attention_head_dim=12, in_channels=16, out_channels=16,
text_dim=32, freq_dim=256, ffn_dim=32, num_layers=2, rope_dim=(4, 4, 4),
),
vae=vae,
scheduler=HeliosScheduler(stage_range=[0, 1], stages=1, use_dynamic_shifting=True),
text_encoder=T5EncoderModel(config),
tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"),
).to("cpu")

kwargs = dict(prompt="dance monkey", negative_prompt="negative", num_inference_steps=1, guidance_scale=1.0, height=16, width=16, num_frames=9, max_sequence_length=16, output_type="latent")
shape = (1, 16, 9, 2, 2)

pipe = make_pipe(); pipe.set_progress_bar_config(disable=True)
a = pipe(**kwargs, generator=torch.Generator("cpu").manual_seed(123), latents=torch.zeros(shape)).frames
pipe = make_pipe(); pipe.set_progress_bar_config(disable=True)
b = pipe(**kwargs, generator=torch.Generator("cpu").manual_seed(123), latents=torch.ones(shape)).frames

print((a - b).abs().max().item()) # 0.0: supplied latents were ignored
```

Relevant precedent:
`FluxPipeline.prepare_latents` uses the provided `latents` when present and validates shape before sampling new noise.

Suggested fix:
```python
if latents is not None:
if isinstance(latents, (list, tuple)):
if len(latents) != num_latent_chunk:
raise ValueError("`latents` must contain one tensor per Helios chunk.")
chunk_latents = latents[k]
elif num_latent_chunk == 1:
chunk_latents = latents
else:
raise ValueError("For multi-chunk Helios generation, pass `latents` as a list of chunk tensors.")
else:
chunk_latents = None

latents = self.prepare_latents(..., latents=chunk_latents)
```

## Issue 3: Non-pyramid Helios passes fractional float timesteps while pyramid and modular cast to `int64`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L815
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py#L950-L952
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/denoise.py#L421-L424

Problem:
`HeliosPipeline` forwards scheduler timesteps directly. With its default sigma path, those timesteps are fractional `float64` values. `HeliosPyramidPipeline` and modular Helios cast the same value to `torch.int64` before calling the transformer.

Impact:
The standard non-pyramid pipeline is numerically inconsistent with the rest of the Helios family and with the model tests, which use integer timesteps. The timestep embedding changes measurably.

Reproduction:
```python
import torch
from diffusers import HeliosTransformer3DModel
from diffusers.utils.torch_utils import randn_tensor

model = HeliosTransformer3DModel(
patch_size=(1, 2, 2), num_attention_heads=2, attention_head_dim=12, in_channels=4, out_channels=4,
text_dim=16, freq_dim=256, ffn_dim=32, num_layers=1, rope_dim=(4, 4, 4),
).eval()
g = torch.Generator("cpu").manual_seed(0)
kwargs = dict(
hidden_states=randn_tensor((1, 4, 2, 16, 16), generator=g),
encoder_hidden_states=randn_tensor((1, 12, 16), generator=g),
indices_hidden_states=torch.ones((1, 2)),
indices_latents_history_short=torch.ones((1, 1)),
indices_latents_history_mid=torch.ones((1, 1)),
indices_latents_history_long=torch.ones((1, 4)),
latents_history_short=randn_tensor((1, 4, 1, 16, 16), generator=g),
latents_history_mid=randn_tensor((1, 4, 1, 16, 16), generator=g),
latents_history_long=randn_tensor((1, 4, 4, 16, 16), generator=g),
return_dict=False,
)
with torch.no_grad():
out_float = model(timestep=torch.tensor([499.5], dtype=torch.float64), **kwargs)[0]
out_int = model(timestep=torch.tensor([499], dtype=torch.int64), **kwargs)[0]
print((out_float - out_int).abs().max().item()) # about 5e-2
```

Relevant precedent:
`HeliosPyramidPipeline` and `HeliosChunkDenoiseInner` both cast timesteps to `torch.int64` before transformer invocation.

Suggested fix:
```python
timestep = t.expand(latents.shape[0]).to(torch.int64)
```

## Issue 4: Precomputed V2V latents hit `UnboundLocalError`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L380-L419
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py#L402-L441

Problem:
`prepare_video_latents` only defines `first_frame_latent` inside `if latents is None`, but returns it unconditionally. Passing precomputed `video_latents` with a raw `video` therefore crashes.

Impact:
The exposed `video_latents` skip-encoding path is unusable unless callers also avoid this helper entirely and provide all derived companion latents themselves.

Reproduction:
```python
from types import SimpleNamespace
import torch
from diffusers import HeliosPipeline

pipe = SimpleNamespace(_execution_device=torch.device("cpu"), vae=SimpleNamespace(dtype=torch.float32), vae_scale_factor_temporal=4)
video = torch.zeros(1, 3, 33, 16, 16)
video_latents = torch.zeros(1, 16, 9, 2, 2)

HeliosPipeline.prepare_video_latents(
pipe,
video=video,
latents_mean=torch.zeros(1, 16, 1, 1, 1),
latents_std=torch.ones(1, 16, 1, 1, 1),
num_latent_frames_per_chunk=9,
dtype=torch.float32,
device=torch.device("cpu"),
latents=video_latents,
)
```

Relevant precedent:
Video-to-video pipelines that accept precomputed latents either validate companion inputs or still derive required first-frame latents from the raw video.

Suggested fix:
```python
first_frame = video[:, :, 0:1, :, :]
first_frame_latent = self.vae.encode(first_frame).latent_dist.sample(generator=generator)
first_frame_latent = (first_frame_latent - latents_mean) * latents_std

if latents is None:
# existing chunk encoding path
...
return first_frame_latent.to(device=device, dtype=dtype), latents.to(device=device, dtype=dtype)
```

Also validate direct `video_latents` calls without raw `video`: require `image_latents` alongside `video_latents`, or raise a clear `ValueError`.

## Issue 5: Prompt cleaning crashes without optional `ftfy`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios.py#L44-L90
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/helios/pipeline_helios_pyramid.py#L45-L103
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/modular_pipelines/helios/encoders.py#L31-L51

Problem:
The modules conditionally import `ftfy`, but `basic_clean` calls `ftfy.fix_text` unconditionally. `ftfy` is not in `install_requires`, so minimal installs can import Helios but fail when encoding any string prompt.

Impact:
Text-to-video generation crashes at prompt encoding in environments that install only core diffusers plus torch/transformers.

Reproduction:
```python
import diffusers.pipelines.helios.pipeline_helios as standard
import diffusers.pipelines.helios.pipeline_helios_pyramid as pyramid
import diffusers.modular_pipelines.helios.encoders as modular

for module in (standard, pyramid, modular):
if hasattr(module, "ftfy"):
delattr(module, "ftfy")
print(module.prompt_clean("hello & world"))
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/wan/pipeline_wan.py#L78-L92

Suggested fix:
```python
def basic_clean(text):
if is_ftfy_available():
text = ftfy.fix_text(text)
text = html.unescape(html.unescape(text))
return text.strip()
```

## Issue 6: Slow/integration coverage is effectively missing

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/helios/test_helios.py#L37-L154
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/helios/test_helios.py#L155-L172
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/modular_pipelines/helios/test_modular_pipeline_helios.py#L127-L166

Problem:
The only standard pipeline slow test is skipped with `TODO: test needs to be implemented`. There is no fast standard test class for `HeliosPyramidPipeline`, and no modular test coverage for `HeliosPyramidDistilledModularPipeline`.

Impact:
The exact public paths with the most branching - pyramid, distilled, slow checkpoint loading, and integration generation - can regress without CI signal. This also hides the runtime issues above.

Reproduction:
```python
from pathlib import Path

standard = Path("tests/pipelines/helios/test_helios.py").read_text()
modular = Path("tests/modular_pipelines/helios/test_modular_pipeline_helios.py").read_text()

print("slow class present:", "class HeliosPipelineIntegrationTests" in standard)
print("slow test skipped:", '@unittest.skip("TODO: test needs to be implemented")' in standard)
print("standard pyramid fast test present:", "HeliosPyramidPipeline" in standard)
print("distilled modular test present:", "HeliosPyramidDistilled" in modular)
```

Relevant precedent:
Existing video pipeline suites generally include a real slow test for the public checkpoint and fast tests for each exported pipeline variant.

Suggested fix:
Add non-skipped slow tests for `HeliosPipeline` and `HeliosPyramidPipeline` using small generation settings, plus fast tests for `HeliosPyramidPipeline` and `HeliosPyramidDistilledModularPipeline`. Add explicit assertions for `latents`, `num_videos_per_prompt`, precomputed I2V/V2V latents, and standard-vs-modular timestep parity.

Note: I attempted to run `tests/pipelines/helios/test_helios.py::HeliosPipelineFastTests::test_inference` with `.venv`, but collection fails in this environment because the installed Torch build lacks `torch._C._distributed_c10d`, imported through shared test utilities.

Start by running the supplied reproductions, then inspect the referenced standard and modular Helios pipeline files around prompt expansion, latent preparation, denoising timesteps, and prepare_video_latents. Compare the non-pyramid, pyramid, and modular paths with the cited Wan and Flux precedents. Done means all four reported behaviors are corrected and covered by regression tests.

Vom Indexierungsmodell aus dem Issue-Text verfasst.