Update app.py

app.py

@@ -215,23 +215,24 @@ def generate_music(
                 Performs model inference to generate music tokens.
                 This function is decorated with @spaces.GPU for GPU usage in Gradio Spaces.
                 """
-                output_seq = model.generate(
-                    input_ids=input_ids,
-                    max_new_tokens=max_new_tokens,
-                    min_new_tokens=100, # Keep min_new_tokens to avoid short generations
-                    do_sample=True,
-                    top_p=top_p,
-                    temperature=temperature,
-                    repetition_penalty=repetition_penalty,
-                    eos_token_id=mmtokenizer.eoa,
-                    pad_token_id=mmtokenizer.eoa,
-                    logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
-                    guidance_scale=guidance_scale,
-                    use_cache=True
-                )
-                if output_seq[0][-1].item() != mmtokenizer.eoa:
-                    tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
-                    output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
+                with torch.inference_mode(), torch.autocast(device_type='cuda', dtype=torch.float16):
+                    output_seq = model.generate(
+                        input_ids=input_ids,
+                        max_new_tokens=max_new_tokens,
+                        min_new_tokens=100, # Keep min_new_tokens to avoid short generations
+                        do_sample=True,
+                        top_p=top_p,
+                        temperature=temperature,
+                        repetition_penalty=repetition_penalty,
+                        eos_token_id=mmtokenizer.eoa,
+                        pad_token_id=mmtokenizer.eoa,
+                        logits_processor=LogitsProcessorList([BlockTokenRangeProcessor(0, 32002), BlockTokenRangeProcessor(32016, 32016)]),
+                        guidance_scale=guidance_scale,
+                        use_cache=True
+                    )
+                    if output_seq[0][-1].item() != mmtokenizer.eoa:
+                        tensor_eoa = torch.as_tensor([[mmtokenizer.eoa]]).to(model.device)
+                        output_seq = torch.cat((output_seq, tensor_eoa), dim=1)
                 return output_seq
     
             output_seq = model_inference(input_ids, max_new_tokens, top_p, temperature, repetition_penalty, guidance_scale)