compatible with L40

Dockerfile

@@ -1,4 +1,4 @@
-FROM witszhang/songgeneration_vllm:v1
+FROM witszhang/songgeneration_vllm:v2
 
 USER root
 

app.py

@@ -183,20 +183,6 @@ lyrics
             )
 
             with gr.Tabs(elem_id="extra-tabs"):
-                with gr.Tab("Genre Select"):
-                    genre = gr.Radio(
-                        choices=["Auto", "Pop", "R&B", "Dance", "Jazz", "Folk", "Rock", "Chinese Style", "Chinese Tradition", "Metal", "Reggae", "Chinese Opera"],
-                        label="Genre Select(Optional)",
-                        value="Auto",
-                        interactive=True,
-                        elem_id="single-select-radio"
-                    )
-                with gr.Tab("Audio Prompt"):
-                    prompt_audio = gr.Audio(
-                        label="Prompt Audio (Optional)",
-                        type="filepath",
-                        elem_id="audio-prompt"
-                    )
                 with gr.Tab("Text Prompt"):
                     gr.Markdown("For detailed usage, please refer to [here](https://github.com/tencent-ailab/SongGeneration?tab=readme-ov-file#-description-input-format)")
                     description = gr.Textbox(
@@ -206,6 +192,20 @@ lyrics
                         lines=1,
                         max_lines=2
                     )
+                with gr.Tab("Audio Prompt"):
+                    prompt_audio = gr.Audio(
+                        label="Prompt Audio (Optional)",
+                        type="filepath",
+                        elem_id="audio-prompt"
+                    )
+                with gr.Tab("Genre Select"):
+                    genre = gr.Radio(
+                        choices=["Auto", "Pop", "R&B", "Dance", "Jazz", "Folk", "Rock", "Chinese Style", "Chinese Tradition", "Metal", "Reggae", "Chinese Opera"],
+                        label="Genre Select(Optional)",
+                        value="Auto",
+                        interactive=True,
+                        elem_id="single-select-radio"
+                    )
 
             with gr.Accordion("Advanced Config", open=False):
                 cfg_coef = gr.Slider(

levo_inference.py

@@ -45,7 +45,7 @@ class LeVoInference(torch.nn.Module):
             model=self.cfg.lm_checkpoint,
             trust_remote_code=True,
             tensor_parallel_size=self.cfg.vllm.device_num,
-            enforce_eager=False,              
+            enforce_eager=True,
             dtype="bfloat16",  
             gpu_memory_utilization=self.cfg.vllm.gpu_memory_utilization, 
             tokenizer=None,

vllm_hacked/v1/sample/sampler.py

@@ -187,7 +187,10 @@ class Sampler(nn.Module):
         # Avoid division by zero if there are greedy requests.
         if not all_random:
             temp = torch.where(temp < _SAMPLING_EPS, 1.0, temp)
-        return logits.div_(temp.unsqueeze(dim=1))
+        try:
+            return logits.div_(temp.view(-1, 1))
+        except:
+            return logits.div_(temp.unsqueeze(dim=1))
 
     def greedy_sample(self, logits: torch.Tensor) -> torch.Tensor:
         return logits.argmax(dim=-1).view(-1)