set inference to false once finished

handler.py

@@ -222,44 +222,6 @@ class EndpointHandler:
             )
         return self.pipe
 
-    def __call__(self, data: Any) -> Dict:
-        """Handle incoming requests."""
-
-        action = data.get("action", None)
-        request_id = data.get("request_id")
-
-        # Check if the request_id is valid for all actions
-        if not request_id:
-            return {"flag": "error", "message": "Missing request_id."}
-
-        if action == "check_progress":
-            return self.check_progress(request_id)
-
-        elif action == "inference":
-            # Check if an inference is already in progress
-            if self.inference_in_progress:
-                return {
-                    "flag": "error",
-                    "message": "Another inference is already in progress. Please wait.",
-                }
-
-            # Set the inference state to in progress
-            self.clean_request_data(request_id)
-            self.inference_in_progress = True
-            self.inference_progress[request_id] = 0
-            self.inference_images[request_id] = None
-
-            self.executor.submit(self.start_inference, data)
-
-            return {
-                "flag": "success",
-                "message": "Inference started",
-                "request_id": request_id,
-            }
-
-        else:
-            return {"flag": "error", "message": f"Unsupported action: {action}"}
-
     def clean_request_data(self, request_id: str):
         """Clean up the data related to a specific request ID."""
 
@@ -418,9 +380,49 @@ class EndpointHandler:
                 num_inference_steps, 0, image, request_id, "complete"
             )
 
+            self.inference_in_progress = False
+
             # for debug
             # image.save("squirelb.png", format="PNG")
 
         except Exception as e:
             # Handle any other exceptions and return an error response
             return {"flag": "error", "message": str(e)}
+
+    def __call__(self, data: Any) -> Dict:
+        """Handle incoming requests."""
+
+        action = data.get("action", None)
+        request_id = data.get("request_id")
+
+        # Check if the request_id is valid for all actions
+        if not request_id:
+            return {"flag": "error", "message": "Missing request_id."}
+
+        if action == "check_progress":
+            return self.check_progress(request_id)
+
+        elif action == "inference":
+            # Check if an inference is already in progress
+            if self.inference_in_progress:
+                return {
+                    "flag": "error",
+                    "message": "Another inference is already in progress. Please wait.",
+                }
+
+            # Set the inference state to in progress
+            self.clean_request_data(request_id)
+            self.inference_in_progress = True
+            self.inference_progress[request_id] = 0
+            self.inference_images[request_id] = None
+
+            self.executor.submit(self.start_inference, data)
+
+            return {
+                "flag": "success",
+                "message": "Inference started",
+                "request_id": request_id,
+            }
+
+        else:
+            return {"flag": "error", "message": f"Unsupported action: {action}"}