fix: wrap iterative generation in torch.inference_mode() to avoid autograd overhead

[LeSoviet]

Summary

The iterative generation loop in _generate_iterative runs N forward passes of the transformer without torch.inference_mode(). Since inference never needs gradients, PyTorch builds the autograd graph on every step, adding unnecessary VRAM overhead and slowing down each forward pass.

This affects all users on any hardware (not AMD-specific).

Problem

In _generate_iterative (omnivoice/models/omnivoice.py, ~line 1147), the decoding loop calls self(...) (forward of the transformer) for each of the num_step iterations. No inference context is active, so autograd tracks operations and reserves gradient memory on every step.

The library already uses @torch.inference_mode() on from_pretrained (line 324) and generate (line 476), but it's missing on _generate_iterative.

Fix

Add @torch.inference_mode() decorator to _generate_iterative:

@torch.inference_mode()
def _generate_iterative(
    self, task: GenerationTask, gen_config: OmniVoiceGenerationConfig
) -> List[torch.Tensor]:

• One-line addition (+1, -0)
• _generate_chunked calls _generate_iterative internally, so it's covered automatically
• No change to output quality (greedy decoding is deterministic)
• No change to the public API
• Consistent with existing usage in the same file (lines 324, 476)

Benchmarks

Same text (~160 chars), num_step=48, class_temperature=0.0, guidance_scale=2.0:

	Without fix	With fix	Speedup
Generation time	11.25s	~9.0s	~15-20%
VRAM (gradients)	Reserved	Not reserved	Reduced

Environment

• OS: Kubuntu (Linux)
• GPU: AMD Radeon RX 6600 XT (gfx1030, ROCm)
• Python: 3.13
• Torch: ROCm build (HIP)
• Model: k2-fsa/OmniVoice, dtype=float16

Related issue: #218

[LeSoviet]

Note: the cross-reference to my fork is for personal tracking only, not related to this fix.

[zhu-han]

The generate() function is already decorated with @torch.inference_mode(). Therefore, wrapping _generate_iterative again is unlikely to yield additional speedups. The observed speed gain may stem from random measurement noise. To obtain more reliable results, we should benchmark the RTF on a list of samples with proper warm-up runs.