[NFC] Remove redundant torch.no_grad in models and pipelines (#854)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -714,7 +714,6 @@ class Bagel(torch.nn.Module):
 
         return generation_input, newlens, new_rope
 
-    @torch.no_grad
     def forward_cache_update_text(
         self,
         past_key_values: NaiveCache,
@@ -821,7 +820,6 @@ class Bagel(torch.nn.Module):
 
         return generation_input, newlens, new_rope
 
-    @torch.no_grad
     def forward_cache_update_vae(
         self,
         vae_model,
@@ -946,7 +944,6 @@ class Bagel(torch.nn.Module):
 
         return generation_input, newlens, new_rope
 
-    @torch.no_grad
     def forward_cache_update_vit(
         self,
         past_key_values: NaiveCache,
@@ -1064,7 +1061,6 @@ class Bagel(torch.nn.Module):
     def prepare_vae_latent(self, curr_kvlens, curr_rope, image_sizes, new_token_ids):
         return self.prepare_input(curr_kvlens, curr_rope, image_sizes, new_token_ids)
 
-    @torch.no_grad
     def generate_image(
         self,
         packed_text_ids: torch.LongTensor,
@@ -1121,7 +1117,6 @@ class Bagel(torch.nn.Module):
         unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
         return unpacked_latent
 
-    @torch.no_grad
     def _forward_flow(
         self,
         x_t: torch.Tensor,

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -644,7 +644,6 @@ class Flux2KleinPipeline(nn.Module, SupportImageInput):
     def interrupt(self):
         return self._interrupt
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py

@@ -516,7 +516,6 @@ class OvisImagePipeline(
     def interrupt(self):
         return self._interrupt
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -347,7 +347,6 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput):
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -278,7 +278,6 @@ class Wan22Pipeline(nn.Module):
     def current_timestep(self):
         return self._current_timestep
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -254,7 +254,6 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput):
         image_embeds = self.image_encoder(pixel_values, output_hidden_states=True)
         return image_embeds.hidden_states[-2]
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -193,7 +193,6 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput):
     def current_timestep(self):
         return self._current_timestep
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -312,7 +312,6 @@ class ZImagePipeline(nn.Module):
     def interrupt(self):
         return self._interrupt
 
-    @torch.no_grad()
     def forward(
         self,
         req: OmniDiffusionRequest,

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py

@@ -1258,7 +1258,6 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
 
         return output
 
-    @torch.no_grad()
     def sample(
         self,
         conditioning_vector,
@@ -1333,7 +1332,6 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
         generated_mel_spectrogram = generated_waveform.permute(0, 2, 1)
         return generated_mel_spectrogram
 
-    @torch.no_grad()
     def fast_block_sample(
         self,
         conditioning_vector: torch.Tensor,