Update app.py

app.py

@@ -21,51 +21,13 @@ 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):
-    
-    headers = {'Content-Type': 'application/json'}
-
-    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
-
-    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!")
-            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}")
+def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
 
-def main(prompt, negative_prompt, subject, number_of_outputs, number_of_images_per_pose, randomise_poses, output_format, output_quality, seed):
     print(f"""
     —/n
-    {prompt}
+    {arg[0]}
     """)
-    is_safe = safety_check(prompt)
+    is_safe = safety_check(arg[0])
     print(is_safe)
     
     match = re.search(r'\bYes\b', is_safe)
@@ -78,7 +40,44 @@ def main(prompt, negative_prompt, subject, number_of_outputs, number_of_images_p
     if status == "Yes" :
         raise gr.Error("Don't ask for such things.")
     else:
-        results = predict(prompt, negative_prompt, subject, number_of_outputs, number_of_images_per_pose, randomise_poses, output_format, output_quality, seed)
+    
+        headers = {'Content-Type': 'application/json'}
+    
+        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
+    
+        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!")
+                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}")
+
 
 title = "Demo for consistent-character cog image by fofr"
 description = "Create images of a given character in different poses • running cog image by fofr"
@@ -155,7 +154,7 @@ with gr.Blocks(css=css) as app:
     outputs = [consistent_results]
 
     submit_btn.click(
-        fn = main,
+        fn = predict,
         inputs = inputs,
         outputs = outputs,
         show_api = False