Update app.py

app.py

@@ -246,10 +246,19 @@ def stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_s
             prompt = prompt[np.newaxis, :]
         print(f"Loaded prompt from {path} with shape: {prompt.shape}")
 
-        # Compute output duration: tokens per second assumed to be 50, only full 6-second segments.
-        output_duration = (prompt.shape[-1] // 50) // 6 * 6
+        # Compute total duration in seconds (assuming 50 tokens per second)
+        total_duration_sec = prompt.shape[-1] // 50
+        if total_duration_sec < 6:
+            # Not enough tokens for a full 6-sec segment; use the entire prompt.
+            output_duration = total_duration_sec
+            print(f"Prompt too short for 6-sec segmentation. Using full duration: {output_duration} seconds.")
+        else:
+            output_duration = (total_duration_sec // 6) * 6
+
+        # If after the above, output_duration is still zero, raise an error.
         if output_duration == 0:
             raise ValueError(f"Output duration computed as 0 for {path}. Prompt length: {prompt.shape[-1]} tokens")
+        
         num_batch = output_duration // 6
 
         # Process prompt in batches
@@ -280,7 +289,7 @@ def stage2_inference(model_stage2, stage1_output_set, stage2_output_dir, batch_s
         # Convert Stage2 output tokens back to numpy using Stage2’s codec manipulator.
         output = codectool_stage2.ids2npy(output)
 
-        # Fix any invalid codes
+        # Fix any invalid codes (if needed)
         fixed_output = copy.deepcopy(output)
         for i, line in enumerate(output):
             for j, element in enumerate(line):