jree423/diffsketcher

1#!/usr/bin/env python2# -*- coding: utf-8 -*-3 4"""5Simplified DiffSketcher implementation for Hugging Face Inference API.6This version doesn't rely on cloning the repository at runtime.7"""8 9import os10import io11import base6412import torch13import numpy as np14from PIL import Image15import cairosvg16import random17from pathlib import Path18 19class SimplifiedDiffSketcher:20    def __init__(self, model_dir):21        """Initialize the simplified DiffSketcher model"""22        self.model_dir = model_dir23        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")24        print(f"Initializing simplified DiffSketcher on device: {self.device}")25        26        # Load CLIP model if available27        try:28            import clip29            self.clip_model, _ = clip.load("ViT-B-32", device=self.device)30            self.clip_available = True31            print("CLIP model loaded successfully")32        except Exception as e:33            print(f"Error loading CLIP model: {e}")34            self.clip_available = False35    36    def generate_svg(self, prompt, num_paths=20, width=512, height=512):37        """Generate an SVG from a text prompt"""38        print(f"Generating SVG for prompt: {prompt}")39        40        # Use CLIP to encode the prompt if available41        if self.clip_available:42            try:43                import clip44                with torch.no_grad():45                    text = clip.tokenize([prompt]).to(self.device)46                    text_features = self.clip_model.encode_text(text)47                    text_features = text_features.cpu().numpy()[0]48                    # Normalize features49                    text_features = text_features / np.linalg.norm(text_features)50            except Exception as e:51                print(f"Error encoding prompt with CLIP: {e}")52                text_features = np.random.randn(512)  # Random features as fallback53        else:54            # Generate random features if CLIP is not available55            text_features = np.random.randn(512)56        57        # Generate a car-like SVG based on the prompt58        svg_content = self._generate_car_svg(prompt, text_features, num_paths, width, height)59        60        return svg_content61    62    def _generate_car_svg(self, prompt, features, num_paths=20, width=512, height=512):63        """Generate a car-like SVG based on the prompt and features"""64        # Start SVG65        svg_content = f"""<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">66            <rect width="100%" height="100%" fill="#f8f8f8"/>67        """68        69        # Use the features to determine car properties70        car_color_hue = int((features[0] + 1) * 180) % 360  # Map to 0-360 hue71        car_size = 0.6 + 0.2 * features[1]  # Size variation72        car_style = int(abs(features[2] * 3)) % 3  # 0: sedan, 1: SUV, 2: sports car73        74        # Calculate car dimensions75        car_width = int(width * 0.7 * car_size)76        car_height = int(height * 0.3 * car_size)77        car_x = (width - car_width) // 278        car_y = height // 279        80        # Generate car body based on style81        if car_style == 0:  # Sedan82            # Car body (rounded rectangle)83            svg_content += f"""<rect x="{car_x}" y="{car_y}" width="{car_width}" height="{car_height}" 84                rx="20" ry="20" fill="hsl({car_color_hue}, 80%, 50%)" stroke="black" stroke-width="2" />"""85            86            # Windshield87            windshield_width = car_width * 0.788            windshield_height = car_height * 0.589            windshield_x = car_x + (car_width - windshield_width) // 290            windshield_y = car_y - windshield_height * 0.391            svg_content += f"""<rect x="{windshield_x}" y="{windshield_y}" width="{windshield_width}" height="{windshield_height}" 92                rx="10" ry="10" fill="#a8d8ff" stroke="black" stroke-width="1" />"""93            94            # Wheels95            wheel_radius = car_height * 0.496            wheel_y = car_y + car_height * 0.897            svg_content += f"""<circle cx="{car_x + car_width * 0.2}" cy="{wheel_y}" r="{wheel_radius}" fill="black" />"""98            svg_content += f"""<circle cx="{car_x + car_width * 0.8}" cy="{wheel_y}" r="{wheel_radius}" fill="black" />"""99            svg_content += f"""<circle cx="{car_x + car_width * 0.2}" cy="{wheel_y}" r="{wheel_radius * 0.6}" fill="#444" />"""100            svg_content += f"""<circle cx="{car_x + car_width * 0.8}" cy="{wheel_y}" r="{wheel_radius * 0.6}" fill="#444" />"""101            102        elif car_style == 1:  # SUV103            # Car body (taller rectangle)104            svg_content += f"""<rect x="{car_x}" y="{car_y - car_height * 0.3}" width="{car_width}" height="{car_height * 1.3}" 105                rx="15" ry="15" fill="hsl({car_color_hue}, 80%, 50%)" stroke="black" stroke-width="2" />"""106            107            # Windshield108            windshield_width = car_width * 0.6109            windshield_height = car_height * 0.6110            windshield_x = car_x + (car_width - windshield_width) // 2111            windshield_y = car_y - car_height * 0.2112            svg_content += f"""<rect x="{windshield_x}" y="{windshield_y}" width="{windshield_width}" height="{windshield_height}" 113                rx="8" ry="8" fill="#a8d8ff" stroke="black" stroke-width="1" />"""114            115            # Wheels (larger)116            wheel_radius = car_height * 0.45117            wheel_y = car_y + car_height * 0.7118            svg_content += f"""<circle cx="{car_x + car_width * 0.2}" cy="{wheel_y}" r="{wheel_radius}" fill="black" />"""119            svg_content += f"""<circle cx="{car_x + car_width * 0.8}" cy="{wheel_y}" r="{wheel_radius}" fill="black" />"""120            svg_content += f"""<circle cx="{car_x + car_width * 0.2}" cy="{wheel_y}" r="{wheel_radius * 0.6}" fill="#444" />"""121            svg_content += f"""<circle cx="{car_x + car_width * 0.8}" cy="{wheel_y}" r="{wheel_radius * 0.6}" fill="#444" />"""122            123        else:  # Sports car124            # Car body (low, sleek shape)125            svg_content += f"""<path d="M {car_x} {car_y + car_height * 0.5} 126                C {car_x + car_width * 0.1} {car_y - car_height * 0.2}, 127                {car_x + car_width * 0.3} {car_y - car_height * 0.3}, 128                {car_x + car_width * 0.5} {car_y - car_height * 0.2} 129                S {car_x + car_width * 0.9} {car_y}, 130                {car_x + car_width} {car_y + car_height * 0.3} 131                L {car_x + car_width} {car_y + car_height * 0.7} 132                C {car_x + car_width * 0.9} {car_y + car_height}, 133                {car_x + car_width * 0.1} {car_y + car_height}, 134                {car_x} {car_y + car_height * 0.7} Z" 135                fill="hsl({car_color_hue}, 90%, 45%)" stroke="black" stroke-width="2" />"""136            137            # Windshield138            windshield_width = car_width * 0.4139            windshield_x = car_x + car_width * 0.3140            windshield_y = car_y - car_height * 0.1141            svg_content += f"""<path d="M {windshield_x} {windshield_y} 142                C {windshield_x + windshield_width * 0.1} {windshield_y - car_height * 0.15}, 143                {windshield_x + windshield_width * 0.9} {windshield_y - car_height * 0.15}, 144                {windshield_x + windshield_width} {windshield_y} Z" 145                fill="#a8d8ff" stroke="black" stroke-width="1" />"""146            147            # Wheels (low profile)148            wheel_radius = car_height * 0.35149            wheel_y = car_y + car_height * 0.7150            svg_content += f"""<ellipse cx="{car_x + car_width * 0.2}" cy="{wheel_y}" rx="{wheel_radius * 1.2}" ry="{wheel_radius}" fill="black" />"""151            svg_content += f"""<ellipse cx="{car_x + car_width * 0.8}" cy="{wheel_y}" rx="{wheel_radius * 1.2}" ry="{wheel_radius}" fill="black" />"""152            svg_content += f"""<ellipse cx="{car_x + car_width * 0.2}" cy="{wheel_y}" rx="{wheel_radius * 0.7}" ry="{wheel_radius * 0.6}" fill="#444" />"""153            svg_content += f"""<ellipse cx="{car_x + car_width * 0.8}" cy="{wheel_y}" rx="{wheel_radius * 0.7}" ry="{wheel_radius * 0.6}" fill="#444" />"""154        155        # Add headlights156        headlight_radius = car_width * 0.05157        headlight_y = car_y + car_height * 0.3158        svg_content += f"""<circle cx="{car_x + car_width * 0.1}" cy="{headlight_y}" r="{headlight_radius}" fill="yellow" stroke="black" stroke-width="1" />"""159        svg_content += f"""<circle cx="{car_x + car_width * 0.9}" cy="{headlight_y}" r="{headlight_radius}" fill="yellow" stroke="black" stroke-width="1" />"""160        161        # Add details based on features162        for i in range(min(10, len(features))):163            feature_val = features[i % len(features)]164            x = car_x + car_width * ((i / 10) * 0.8 + 0.1)165            y = car_y + car_height * ((feature_val + 1) / 4)166            size = car_width * 0.03 * abs(feature_val)167            svg_content += f"""<circle cx="{x}" cy="{y}" r="{size}" fill="rgba(0,0,0,0.2)" />"""168        169        # Add prompt as text170        svg_content += f"""<text x="{width/2}" y="{height - 20}" font-family="Arial" font-size="12" text-anchor="middle">{prompt}</text>"""171        172        # Close SVG173        svg_content += "</svg>"174        175        return svg_content176    177    def svg_to_png(self, svg_content):178        """Convert SVG content to PNG"""179        try:180            png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))181            return png_data182        except Exception as e:183            print(f"Error converting SVG to PNG: {e}")184            # Create a simple error image185            image = Image.new("RGB", (512, 512), color="#ff0000")186            from PIL import ImageDraw187            draw = ImageDraw.Draw(image)188            draw.text((256, 256), f"Error: {str(e)}", fill="white", anchor="mm")189            190            # Convert PIL Image to PNG data191            buffer = io.BytesIO()192            image.save(buffer, format="PNG")193            return buffer.getvalue()194    195    def __call__(self, prompt):196        """Generate an SVG from a text prompt and convert to PNG"""197        svg_content = self.generate_svg(prompt)198        png_data = self.svg_to_png(svg_content)199        200        # Create a PIL Image from the PNG data201        image = Image.open(io.BytesIO(png_data))202        203        # Create the response204        response = {205            "svg": svg_content,206            "svg_base64": base64.b64encode(svg_content.encode("utf-8")).decode("utf-8"),207            "png_base64": base64.b64encode(png_data).decode("utf-8"),208            "image": image209        }210        211        return response