Upload folder using huggingface_hub

handler.py

@@ -0,0 +1,70 @@
+from typing import Dict, List, Any
+import torch
+import base64
+import io
+from PIL import Image
+from diffusers import DiffusionPipeline
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Initialize the model
+        self.model = DiffusionPipeline.from_pretrained(path)
+        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+        
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        """
+        Args:
+            data: Dictionary with the following structure:
+                {
+                    "prompt": str,
+                    "negative_prompt": str (optional),
+                    "num_paths": int (optional),
+                    "num_iter": int (optional),
+                    "guidance_scale": float (optional),
+                    "width": int (optional),
+                    "seed": int (optional)
+                }
+        Returns:
+            Dictionary with the following structure:
+                {
+                    "svg": str,
+                    "image": str (base64 encoded image)
+                }
+        """
+        # Extract parameters from the input data
+        prompt = data.get("prompt", "")
+        negative_prompt = data.get("negative_prompt", None)
+        num_paths = data.get("num_paths", 96)
+        num_iter = data.get("num_iter", 800)
+        guidance_scale = data.get("guidance_scale", 7.5)
+        width = data.get("width", 2)
+        seed = data.get("seed", None)
+        
+        # Set the seed if provided
+        if seed is not None:
+            torch.manual_seed(seed)
+            
+        # Generate the SVG
+        output = self.model(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_paths=num_paths,
+            num_iter=num_iter,
+            guidance_scale=guidance_scale,
+            width=width
+        )
+        
+        # Get the SVG and image
+        svg = output.svg
+        image = output.images[0]
+        
+        # Convert the image to base64
+        buffered = io.BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode()
+        
+        # Return the results
+        return {
+            "svg": svg,
+            "image": img_str
+        }