Update app.py

app.py

@@ -22,62 +22,67 @@ from utils.gradio_helpers import parse_outputs, process_outputs
 names = ['prompt', 'negative_prompt', 'subject', 'number_of_outputs', 'number_of_images_per_pose', 'randomise_poses', 'output_format', 'output_quality', 'seed']
 
 def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
-
     print(f"""
-    —/n
-    {args[0]}
-    """)
-    is_safe = safety_check(args[0])
-    print(is_safe)
+        —/n
+        {args[0]}
+        """)
+    
+    try:
+        
+        is_safe = safety_check(args[0])
+        print(is_safe)
     
-    match = re.search(r'\bYes\b', is_safe)
+        match = re.search(r'\bYes\b', is_safe)
     
-    if match:
-        status = 'Yes'
-    else:
-        status = None
+        if match:
+            status = 'Yes'
+        else:
+            status = None
     
-    if status == "Yes" :
-        raise gr.Error("Don't ask for such things.")
-    else:
+        if status == "Yes" :
+            raise gr.Error("Don't ask for such things.")
+        else:
     
-        headers = {'Content-Type': 'application/json'}
+            headers = {'Content-Type': 'application/json'}
     
-        payload = {"input": {}}
+            payload = {"input": {}}
         
         
-        base_url = "http://0.0.0.0:7860"
-        for i, key in enumerate(names):
-            value = args[i]
-            if value and (os.path.exists(str(value))):
-                value = f"{base_url}/file=" + value
-            if value is not None and value != "":
-                payload["input"][key] = value
+            base_url = "http://0.0.0.0:7860"
+            for i, key in enumerate(names):
+                value = args[i]
+                if value and (os.path.exists(str(value))):
+                    value = f"{base_url}/file=" + value
+                if value is not None and value != "":
+                    payload["input"][key] = value
     
-        response = requests.post("http://0.0.0.0:5000/predictions", headers=headers, json=payload)
+            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)
-            while response.json()["status"] != "succeeded":
-                if response.json()["status"] == "failed":
-                    raise gr.Error("The submission failed!")
+            if response.status_code == 201:
+                follow_up_url = response.json()["urls"]["get"]
                 response = requests.get(follow_up_url, headers=headers)
-                time.sleep(1)
-        if response.status_code == 200:
-            json_response = response.json()
-            #If the output component is JSON return the entire output response 
-            if(outputs[0].get_config()["name"] == "json"):
-                return json_response["output"]
-            predict_outputs = parse_outputs(json_response["output"])
-            processed_outputs = process_outputs(predict_outputs)        
-            return tuple(processed_outputs) if len(processed_outputs) > 1 else processed_outputs[0]
-        else:
-            if(response.status_code == 409):
-                raise gr.Error(f"Sorry, the Cog image is still processing. Try again in a bit.")
-            raise gr.Error(f"The submission failed! Error: {response.status_code}")
-
+                while response.json()["status"] != "succeeded":
+                    if response.json()["status"] == "failed":
+                        raise gr.Error("The submission failed!")
+                    response = requests.get(follow_up_url, headers=headers)
+                    time.sleep(1)
+            if response.status_code == 200:
+                json_response = response.json()
+                #If the output component is JSON return the entire output response 
+                if(outputs[0].get_config()["name"] == "json"):
+                    return json_response["output"]
+                predict_outputs = parse_outputs(json_response["output"])
+                processed_outputs = process_outputs(predict_outputs)        
+                return tuple(processed_outputs) if len(processed_outputs) > 1 else processed_outputs[0]
+            else:
+                if(response.status_code == 409):
+                    raise gr.Error(f"Sorry, the Cog image is still processing. Try again in a bit.")
+                raise gr.Error(f"The submission failed! Error: {response.status_code}")
+
+    except Exception as e:
+        # Handle any other type of error
+        raise gr.Error(f"An error occurred: {e}")
 
 title = "Demo for consistent-character cog image by fofr"
 description = "Create images of a given character in different poses • running cog image by fofr"