Update app.py

app.py

@@ -5,12 +5,53 @@ import time
 import os
 import spaces
 import torch
+import subprocess
+import signal
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device)  # <-- 'cpu' 🤔
 
 names = ['prompt', 'negative_prompt', 'subject', 'number_of_outputs', 'number_of_images_per_pose', 'randomise_poses', 'output_format', 'output_quality', 'seed']
 
+def check_cog_server():
+    try:
+        # Start the Cog server in the background
+        cog_process = subprocess.Popen(["python3", "-m", "cog.server.http", "--threads=10"], cwd="/src")
+
+        # Wait for the Cog server to start on port 5000
+        counter1 = 0
+        while True:
+            try:
+                requests.get("http://localhost:5000")
+                print("Cog server is running on port 5000.")
+                break
+            except requests.exceptions.ConnectionError:
+                print("Waiting for Cog server to start on port 5000...")
+                time.sleep(5)
+                counter1 += 1
+                if counter1 >= 250:
+                    raise Exception("Error: Cog server did not start on port 5000 after 250 attempts.")
+
+        # Wait for the Cog server to be fully ready
+        counter2 = 0
+        while True:
+            response = requests.get("http://localhost:5000/health-check")
+            status = response.json().get("status")
+            if status == "READY":
+                print("Cog server is fully ready.")
+                break
+            else:
+                print("Waiting for Cog server (models loading) on port 5000...")
+                time.sleep(5)
+                counter2 += 1
+                if counter2 >= 250:
+                    raise Exception("Error: Cog server did not become fully ready after 250 attempts.")
+
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        cog_process.send_signal(signal.SIGINT)  # Send interrupt signal to the Cog process
+        raise e
+
 @spaces.GPU
 def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
     print(zero.device)  # <-- 'cuda:0' 🤗
@@ -18,7 +59,6 @@ def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
 
     payload = {"input": {}}
     
-    
     base_url = "http://0.0.0.0:7860"
     for i, key in enumerate(names):
         value = args[i]
@@ -29,7 +69,6 @@ def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
 
     response = requests.post("http://0.0.0.0:5000/predictions", headers=headers, json=payload)
 
-    
     if response.status_code == 201:
         follow_up_url = response.json()["urls"]["get"]
         response = requests.get(follow_up_url, headers=headers)
@@ -132,4 +171,7 @@ with gr.Blocks(css=css) as app:
         show_api = False
     )
 
+# Check the Cog server's readiness before launching the Gradio app
+check_cog_server()
+
 app.queue(max_size=12, api_open=False).launch(share=False, show_api=False)