Autorestart space: Cuda Error (#15)

- Autorestart space (024283f748e35613e326c48c37d82ea8637f89d4)


Co-authored-by: Max Skobeev <[email protected]>

app.py

@@ -119,90 +119,98 @@ def build_html_error_message(error):
 @GPU_DECORATOR
 @torch.inference_mode()
 def inference(req: ServeTTSRequest):
-    # Parse reference audio aka prompt
-    refs = req.references
-
-    prompt_tokens = [
-        encode_reference(
-            decoder_model=decoder_model,
-            reference_audio=ref.audio,
-            enable_reference_audio=True,
-        )
-        for ref in refs
-    ]
-    prompt_texts = [ref.text for ref in refs]
-
-    if req.seed is not None:
-        set_seed(req.seed)
-        logger.warning(f"set seed: {req.seed}")
-
-    # LLAMA Inference
-    request = dict(
-        device=decoder_model.device,
-        max_new_tokens=req.max_new_tokens,
-        text=(
-            req.text
-            if not req.normalize
-            else ChnNormedText(raw_text=req.text).normalize()
-        ),
-        top_p=req.top_p,
-        repetition_penalty=req.repetition_penalty,
-        temperature=req.temperature,
-        compile=args.compile,
-        iterative_prompt=req.chunk_length > 0,
-        chunk_length=req.chunk_length,
-        max_length=4096,
-        prompt_tokens=prompt_tokens,
-        prompt_text=prompt_texts,
-    )
-
-    response_queue = queue.Queue()
-    llama_queue.put(
-        GenerateRequest(
-            request=request,
-            response_queue=response_queue,
-        )
-    )
-
-    segments = []
-
-    while True:
-        result: WrappedGenerateResponse = response_queue.get()
-        if result.status == "error":
-            yield None, None, build_html_error_message(result.response)
-            break
-
-        result: GenerateResponse = result.response
-        if result.action == "next":
-            break
-
-        with autocast_exclude_mps(
-            device_type=decoder_model.device.type, dtype=args.precision
-        ):
-            fake_audios = decode_vq_tokens(
+    try:
+        # Parse reference audio aka prompt
+        refs = req.references
+    
+        prompt_tokens = [
+            encode_reference(
                 decoder_model=decoder_model,
-                codes=result.codes,
+                reference_audio=ref.audio,
+                enable_reference_audio=True,
             )
-
-        fake_audios = fake_audios.float().cpu().numpy()
-        segments.append(fake_audios)
-
-    if len(segments) == 0:
-        return (
-            None,
-            None,
-            build_html_error_message(
-                i18n("No audio generated, please check the input text.")
+            for ref in refs
+        ]
+        prompt_texts = [ref.text for ref in refs]
+    
+        if req.seed is not None:
+            set_seed(req.seed)
+            logger.warning(f"set seed: {req.seed}")
+    
+        # LLAMA Inference
+        request = dict(
+            device=decoder_model.device,
+            max_new_tokens=req.max_new_tokens,
+            text=(
+                req.text
+                if not req.normalize
+                else ChnNormedText(raw_text=req.text).normalize()
             ),
+            top_p=req.top_p,
+            repetition_penalty=req.repetition_penalty,
+            temperature=req.temperature,
+            compile=args.compile,
+            iterative_prompt=req.chunk_length > 0,
+            chunk_length=req.chunk_length,
+            max_length=4096,
+            prompt_tokens=prompt_tokens,
+            prompt_text=prompt_texts,
         )
-
-    # No matter streaming or not, we need to return the final audio
-    audio = np.concatenate(segments, axis=0)
-    yield None, (decoder_model.spec_transform.sample_rate, audio), None
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        gc.collect()
+    
+        response_queue = queue.Queue()
+        llama_queue.put(
+            GenerateRequest(
+                request=request,
+                response_queue=response_queue,
+            )
+        )
+    
+        segments = []
+    
+        while True:
+            result: WrappedGenerateResponse = response_queue.get()
+            if result.status == "error":
+                yield None, None, build_html_error_message(result.response)
+                break
+    
+            result: GenerateResponse = result.response
+            if result.action == "next":
+                break
+    
+            with autocast_exclude_mps(
+                device_type=decoder_model.device.type, dtype=args.precision
+            ):
+                fake_audios = decode_vq_tokens(
+                    decoder_model=decoder_model,
+                    codes=result.codes,
+                )
+    
+            fake_audios = fake_audios.float().cpu().numpy()
+            segments.append(fake_audios)
+    
+        if len(segments) == 0:
+            return (
+                None,
+                None,
+                build_html_error_message(
+                    i18n("No audio generated, please check the input text.")
+                ),
+            )
+    
+        # No matter streaming or not, we need to return the final audio
+        audio = np.concatenate(segments, axis=0)
+        yield None, (decoder_model.spec_transform.sample_rate, audio), None
+    
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+    except Exception as e:
+        er = "CUDA error: device-side assert triggered"
+        if er in e:
+            app.close()
+        else:
+            raise Exception(e)
 
 n_audios = 4