Update app.py

app.py

@@ -109,8 +109,20 @@ def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalt
 
 @spaces.GPU
 def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
-    if model_manager.model is not None:
-        model_manager.model = model_manager.model.to("cuda")
+    max_retries = 5
+        for attempt in range(max_retries):
+            try:
+                if model_manager.model is not None:
+                    model_manager.model = model_manager.model.to("cuda")
+                break
+            except RuntimeError as e:
+                if "CUDA" in str(e) and attempt < max_retries - 1:
+                    yield f"Waiting for Hugging Face ZeroGPU allocation (Attempt {attempt+1}/{max_retries})...", "Queueing..."
+                    time.sleep(2)
+                else:
+                    yield f"ZeroGPU initialization failed: {str(e)}. Try clicking Generate again.", "GPU Unavailable"
+                    return
+
     yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=True)
 
 def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=False):