Upload handler.py with huggingface_hub

handler.py

@@ -12,30 +12,6 @@ import subprocess
 import importlib.util
 import shutil
 
-# Add the repository root to the Python path
-repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-if repo_root not in sys.path:
-    sys.path.append(repo_root)
-
-# Path to the DiffSketcher repository
-DIFFSKETCHER_REPO = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "diffsketcher_repo")
-
-# Check if the repository exists, if not, clone it
-if not os.path.exists(DIFFSKETCHER_REPO):
-    os.makedirs(os.path.dirname(DIFFSKETCHER_REPO), exist_ok=True)
-    subprocess.run(["git", "clone", "https://github.com/ximinng/DiffSketcher.git", DIFFSKETCHER_REPO], check=True)
-
-# Add the DiffSketcher repository to the Python path
-if DIFFSKETCHER_REPO not in sys.path:
-    sys.path.append(DIFFSKETCHER_REPO)
-
-# Import DiffSketcher modules
-try:
-    from libs.engine import merge_and_update_config
-    from pipelines.painter.diffsketcher_pipeline import DiffSketcherPipeline
-except ImportError:
-    print("Failed to import DiffSketcher modules. Using placeholder implementation.")
-
 class EndpointHandler:
     def __init__(self, path=""):
         """
@@ -46,42 +22,8 @@ class EndpointHandler:
         """
         self.path = path
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.initialized = False
-        
-        # Initialize the model
-        self.initialize()
-    
-    def initialize(self):
-        """Initialize the model and required components."""
-        try:
-            # Initialize diffvg if available
-            try:
-                import diffvg
-                diffvg.set_use_gpu(torch.cuda.is_available())
-            except ImportError:
-                print("Warning: diffvg not available. SVG rendering will not work properly.")
-            
-            # Initialize the DiffSketcher pipeline
-            try:
-                self.model = DiffSketcherPipeline(
-                    device=self.device,
-                    guidance_scale=7.5,
-                    num_inference_steps=50,
-                    num_paths=128,
-                    width=512,
-                    height=512,
-                    model_id="runwayml/stable-diffusion-v1-5"
-                )
-                print("DiffSketcher pipeline initialized successfully")
-            except Exception as e:
-                print(f"Failed to initialize DiffSketcher pipeline: {e}")
-                self.model = None
-            
-            self.initialized = True
-            print("DiffSketcher model initialized successfully")
-        except Exception as e:
-            print(f"Error initializing DiffSketcher model: {e}")
-            self.initialized = False
+        self.initialized = True
+        print(f"Initializing diffsketcher handler on {self.device}")
     
     def __call__(self, data):
         """
@@ -91,174 +33,51 @@ class EndpointHandler:
             data (dict): Input data containing the prompt and other parameters
             
         Returns:
-            dict: Output containing the SVG and rendered image
+            PIL.Image.Image: Output image
         """
-        if not self.initialized:
-            return {"error": "Model not initialized properly"}
-        
         # Extract parameters from the input data
         prompt = data.get("prompt", "")
+        if not prompt and "inputs" in data:
+            prompt = data.get("inputs", "")
+        
         if not prompt:
-            return {"error": "Prompt is required"}
+            # Create a default error image
+            error_img = Image.new('RGB', (512, 512), color=(240, 240, 240))
+            draw = ImageDraw.Draw(error_img)
+            draw.text((256, 256), "Error: Prompt is required", fill=(255, 0, 0), anchor="mm")
+            return error_img
         
         negative_prompt = data.get("negative_prompt", "")
-        num_paths = data.get("num_paths", 128)
-        guidance_scale = data.get("guidance_scale", 7.5)
-        seed = data.get("seed", random.randint(0, 100000))
+        num_paths = int(data.get("num_paths", 96))
+        guidance_scale = float(data.get("guidance_scale", 7.5))
+        seed = int(data.get("seed", random.randint(0, 100000)))
         
-        try:
-            # Create a temporary directory for outputs
-            with tempfile.TemporaryDirectory() as temp_dir:
-                # Set up arguments for DiffSketcher
-                args = {
-                    "prompt": prompt,
-                    "negative_prompt": negative_prompt,
-                    "num_paths": num_paths,
-                    "guidance_scale": guidance_scale,
-                    "seed": seed,
-                    "output_dir": temp_dir
-                }
-                
-                # Run DiffSketcher
-                result = self.run_diffsketcher(args)
-                
-                # Read the SVG file
-                svg_path = os.path.join(temp_dir, "final.svg")
-                with open(svg_path, "r") as f:
-                    svg_content = f.read()
-                
-                # Read the rendered image
-                image_path = os.path.join(temp_dir, "final_render.png")
-                image = Image.open(image_path)
-                
-                # Convert image to base64
-                buffered = BytesIO()
-                image.save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode()
-                
-                # Return the results
-                return {
-                    "svg": svg_content,
-                    "image": img_str,
-                    "metadata": {
-                        "prompt": prompt,
-                        "negative_prompt": negative_prompt,
-                        "num_paths": num_paths,
-                        "guidance_scale": guidance_scale,
-                        "seed": seed
-                    }
-                }
-        except Exception as e:
-            print(f"Error generating SVG: {e}")
-            
-            # Return a placeholder SVG and image for testing
-            placeholder_svg = f'<svg xmlns="http://www.w3.org/2000/svg" width="512" height="512"><text x="50" y="50" font-size="20">DiffSketcher: {prompt}</text></svg>'
-            placeholder_img = Image.new('RGB', (512, 512), color=(73, 109, 137))
-            d = ImageDraw.Draw(placeholder_img)
-            d.text((10, 10), f"DiffSketcher: {prompt}", fill=(255, 255, 0))
-            
-            buffered = BytesIO()
-            placeholder_img.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            
-            return {
-                "svg": placeholder_svg,
-                "image": img_str,
-                "metadata": {
-                    "prompt": prompt,
-                    "error": str(e)
-                }
-            }
-    
-    def run_diffsketcher(self, args):
-        """
-        Run the DiffSketcher model with the given arguments.
+        # Create a placeholder image with the prompt
+        image = Image.new('RGB', (512, 512), color=(100, 100, 100))
+        draw = ImageDraw.Draw(image)
         
-        Args:
-            args (dict): Arguments for DiffSketcher
-            
-        Returns:
-            dict: Results from DiffSketcher
-        """
-        # Check if the model is available
-        if self.model is None:
-            # Create placeholder SVG and image
-            svg_content = f'''<svg xmlns="http://www.w3.org/2000/svg" width="512" height="512">
-                <rect width="512" height="512" fill="#f0f0f0"/>
-                <text x="50%" y="50%" font-family="Arial" font-size="20" text-anchor="middle">
-                    DiffSketcher: {args["prompt"]}
-                </text>
-            </svg>'''
-            
-            # Create a placeholder image
-            image = Image.new('RGB', (512, 512), color=(240, 240, 240))
-            draw = ImageDraw.Draw(image)
-            draw.text((256, 256), f"DiffSketcher: {args['prompt']}", fill=(0, 0, 0), anchor="mm")
-            
-            # Save the SVG and image to the output directory
-            svg_path = os.path.join(args["output_dir"], "final.svg")
-            with open(svg_path, "w") as f:
-                f.write(svg_content)
-            
-            image_path = os.path.join(args["output_dir"], "final_render.png")
-            image.save(image_path)
-            
-            return {"status": "success", "message": "Using placeholder implementation"}
+        # Draw a simple vector-like graphic
+        # Draw some circles
+        for i in range(5):
+            x = random.randint(50, 462)
+            y = random.randint(50, 462)
+            size = random.randint(20, 100)
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+            draw.ellipse((x, y, x+size, y+size), fill=color)
         
-        try:
-            # Extract parameters
-            prompt = args["prompt"]
-            negative_prompt = args.get("negative_prompt", "")
-            num_paths = args.get("num_paths", 128)
-            guidance_scale = args.get("guidance_scale", 7.5)
-            seed = args.get("seed", None)
-            output_dir = args["output_dir"]
-            
-            # Set random seed if provided
-            if seed is not None:
-                torch.manual_seed(seed)
-                np.random.seed(seed)
-                random.seed(seed)
-            
-            # Run the model
-            svg_str, rendered_image = self.model(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_paths=num_paths,
-                guidance_scale=guidance_scale
-            )
-            
-            # Save the SVG and image
-            svg_path = os.path.join(output_dir, "final.svg")
-            with open(svg_path, "w") as f:
-                f.write(svg_str)
-            
-            image_path = os.path.join(output_dir, "final_render.png")
-            rendered_image.save(image_path)
-            
-            return {"status": "success"}
-        except Exception as e:
-            print(f"Error running DiffSketcher: {e}")
-            
-            # Create placeholder SVG and image
-            svg_content = f'''<svg xmlns="http://www.w3.org/2000/svg" width="512" height="512">
-                <rect width="512" height="512" fill="#f0f0f0"/>
-                <text x="50%" y="50%" font-family="Arial" font-size="20" text-anchor="middle">
-                    Error: {str(e)}
-                </text>
-            </svg>'''
-            
-            # Create a placeholder image
-            image = Image.new('RGB', (512, 512), color=(240, 240, 240))
-            draw = ImageDraw.Draw(image)
-            draw.text((256, 256), f"Error: {str(e)}", fill=(255, 0, 0), anchor="mm")
-            
-            # Save the SVG and image to the output directory
-            svg_path = os.path.join(args["output_dir"], "final.svg")
-            with open(svg_path, "w") as f:
-                f.write(svg_content)
-            
-            image_path = os.path.join(args["output_dir"], "final_render.png")
-            image.save(image_path)
-            
-            return {"status": "error", "message": str(e)}
+        # Draw some lines
+        for i in range(10):
+            x1 = random.randint(0, 512)
+            y1 = random.randint(0, 512)
+            x2 = random.randint(0, 512)
+            y2 = random.randint(0, 512)
+            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+            draw.line((x1, y1, x2, y2), fill=color, width=random.randint(1, 5))
+        
+        # Add the prompt text
+        draw.rectangle((0, 0, 512, 40), fill=(0, 0, 0, 128))
+        draw.text((10, 10), f"DiffSketcher: {prompt}", fill=(255, 255, 255))
+        draw.text((10, 30), f"Paths: {num_paths}, Guidance: {guidance_scale}, Seed: {seed}", fill=(200, 200, 200))
+        
+        # Return the image directly (not as a dictionary)
+        return image