Upload 3 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+image.png filter=lfs diff=lfs merge=lfs -text

check.py

@@ -0,0 +1,8 @@
+from handler import EndpointHandler
+
+my_handler = EndpointHandler()
+payload = {"inputs": {"prompt": "a dog waiting for its companion to come."}}
+
+# test the handler
+image = my_handler(payload)
+image.save("image.png")

handler.py

@@ -0,0 +1,25 @@
+from diffusers import DiffusionPipeline 
+from typing import Any, Dict, List
+import torch 
+
+class EndpointHandler:
+    def __init__(self):
+        self.pipeline = DiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16
+        ).to("cuda")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs")
+        prompt = inputs.pop("prompt", "dog")
+        guidance_scale = inputs.pop("guidance_scale", 4.5)
+        num_inference_steps = inputs.pop("num_inference_steps", 28)
+        seed = inputs.pop("seed", 0)
+        
+        prediction = self.pipeline(
+            prompt, 
+            negative_prompt="bad quality, worse quality, deformed",
+            guidance_scale=guidance_scale, 
+            num_inference_steps=num_inference_steps, 
+            generator=torch.manual_seed(seed)
+        ).images[0]
+        return prediction