New handler

handler.py

@@ -87,9 +87,7 @@ class EndpointHandler:
         self.total_steps = {}
         self.inference_in_progress = False
 
-        self.executor = ThreadPoolExecutor(
-            max_workers=1
-        )  # Vous pouvez ajuster max_workers en fonction de vos besoins
+        self.executor = ThreadPoolExecutor(max_workers=1)
 
         # load the optimized model
         self.pipe = DiffusionPipeline.from_pretrained(
@@ -226,13 +224,13 @@ class EndpointHandler:
         """Clean up the data related to a specific request ID."""
 
         # Remove the request ID from the progress dictionary
-        self.inference_progress.pop(request_id, None)
+        self.inference_progress.clear()
 
         # Remove the request ID from the images dictionary
-        self.inference_images.pop(request_id, None)
+        self.inference_images.clear()
 
         # Remove the request ID from the total_steps dictionary
-        self.total_steps.pop(request_id, None)
+        self.total_steps.clear()
 
         # Set inference to False
         self.inference_in_progress = False
@@ -269,6 +267,7 @@ class EndpointHandler:
 
         except Exception as e:
             print(f"Error: {e}")
+            raise
 
         # Store progress and image
         progress_percentage = (
@@ -278,7 +277,9 @@ class EndpointHandler:
         self.inference_progress[request_id] = progress_percentage
         self.inference_images[request_id] = img_str
 
-    def check_progress(self, request_id: str) -> Dict[str, Union[str, float]]:
+        print(self.inference_progress)
+
+    def check_progress(self, request_id: str) -> dict[str, str | float]:
         progress = self.inference_progress.get(request_id, 0)
         latest_image = self.inference_images.get(request_id, None)
 
@@ -387,7 +388,8 @@ class EndpointHandler:
 
         except Exception as e:
             # Handle any other exceptions and return an error response
-            return {"flag": "error", "message": str(e)}
+            print(f"Error in start inference {e}")
+            raise
 
     def __call__(self, data: Any) -> Dict:
         """Handle incoming requests."""