Upload handler.py with huggingface_hub

handler.py

@@ -1,16 +1,11 @@
-from typing import Dict, List, Any
+
+from typing import Dict, Any
 import torch
 import base64
 import io
 import os
 import json
 from PIL import Image
-from diffusers import DiffusionPipeline
-
-class DiffSketcherPipeline:
-    def __init__(self):
-        # This is a placeholder class that will be replaced by the actual implementation
-        pass
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -22,7 +17,7 @@ class EndpointHandler:
         else:
             # Create a default config
             self.config = {
-                "architecture": "DiffSketcherPipeline",
+                "architecture": "SimplePipeline",
                 "format": "diffusers",
                 "version": "0.1.0"
             }
@@ -30,48 +25,22 @@ class EndpointHandler:
             with open(model_index_path, "w") as f:
                 json.dump(self.config, f, indent=2)
         
-        # Initialize a simple pipeline for now
-        self.model = DiffSketcherPipeline()
+        # Initialize device
         self.device = 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
+        # Extract prompt 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)
+        if not prompt and "prompts" in data:
+            prompts = data.get("prompts", [""])
+            prompt = prompts[0] if prompts else ""
         
-        # Set the seed if provided
-        if seed is not None:
-            torch.manual_seed(seed)
-            
         # Generate a placeholder SVG
-        svg = f'<svg xmlns="http://www.w3.org/2000/svg" width="512" height="512" viewBox="0 0 512 512"><text x="50%" y="50%" dominant-baseline="middle" text-anchor="middle" font-size="20">DiffSketcher: {prompt}</text></svg>'
+        model_name = os.path.basename(os.getcwd())
+        svg = f'<svg xmlns="http://www.w3.org/2000/svg" width="512" height="512" viewBox="0 0 512 512"><text x="50%" y="50%" dominant-baseline="middle" text-anchor="middle" font-size="20">{model_name}: {prompt}</text></svg>'
         
         # Create a placeholder image
-        image = Image.new('RGB', (512, 512), color = (73, 109, 137))
+        image = Image.new('RGB', (512, 512), color = (100, 100, 100))
         
         # Convert the image to base64
         buffered = io.BytesIO()
@@ -82,4 +51,5 @@ class EndpointHandler:
         return {
             "svg": svg,
             "image": img_str
-        }
+        }
+