Update app.py

app.py

@@ -206,7 +206,7 @@ def stage2_generate(model_stage2, prompt, batch_size=16):
     for frames_idx in range(codec_ids_tensor.shape[1]):
         cb0 = codec_ids_tensor[:, frames_idx:frames_idx+1]
         prompt_ids_tensor = torch.cat([prompt_ids_tensor, cb0], dim=1)
-        with torch.no_grad():
+        with torch.inference_mode(), torch.cuda.amp.autocast(dtype=torch.float16):
             stage2_output = model_stage2.generate(
                 input_ids=prompt_ids_tensor,
                 min_new_tokens=7,
@@ -214,6 +214,7 @@ def stage2_generate(model_stage2, prompt, batch_size=16):
                 eos_token_id=mmtokenizer.eoa,
                 pad_token_id=mmtokenizer.eoa,
                 logits_processor=block_list,
+                use_cache=True
             )
         # Ensure exactly 7 new tokens were added.
         assert stage2_output.shape[1] - prompt_ids_tensor.shape[1] == 7, (