Update handler.py

handler.py

@@ -9,16 +9,26 @@ class EndpointHandler:
         ).to("cuda")
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        data = data.pop("json")
-        prompt = data.pop("inputs", None)
-        parameters = data.pop("parameters", {})
-        negative_prompt = parameters.pop("negative_prompt", "bad quality, worse quality, deformed")
-        height = parameters.pop("height", 512)
-        width = parameters.pop("width", 512)
-        guidance_scale = parameters.pop("guidance_scale", 4.5)
-        num_inference_steps = parameters.pop("num_inference_steps", 28)
-        seed = parameters.pop("seed", 0)
-        
+        # Extract data
+        data = data.get("json", data)
+        prompt = data.get("inputs", None)
+        parameters = data.get("parameters", {})
+        if not prompt:
+            raise ValueError("Input prompt is missing.")
+
+
+        # Extract parameters with defaults
+        negative_prompt = parameters.get("negative_prompt", "bad quality, worse quality, deformed")
+        height = parameters.get("height", 512)
+        width = parameters.get("width", 512)
+        guidance_scale = parameters.get("guidance_scale", 4.5)
+        num_inference_steps = parameters.get("num_inference_steps", 28)
+        seed = parameters.get("seed", 0)
+
+        # Seed generator
+        generator = torch.manual_seed(seed)
+
+        # Generate prediction
         prediction = self.pipeline(
             prompt, 
             negative_prompt=negative_prompt,
@@ -26,6 +36,6 @@ class EndpointHandler:
             width=width,
             guidance_scale=guidance_scale, 
             num_inference_steps=num_inference_steps, 
-            generator=torch.manual_seed(seed)
+            generator=generator
         ).images[0]
         return prediction