jree423/diffsketcher

1#!/usr/bin/env python2# -*- coding: utf-8 -*-3 4"""5Simplified DiffSketcher model for text-to-SVG generation.6"""7 8import os9import io10import base6411import torch12import numpy as np13from PIL import Image14import clip15import torch.nn.functional as F16import xml.etree.ElementTree as ET17import cairosvg18 19class DiffSketcherModel:20    def __init__(self, model_dir):21        """Initialize the DiffSketcher model"""22        self.model_dir = model_dir23        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")24        25        # Load CLIP model26        self.clip_model_path = os.path.join(model_dir, "ViT-B-32.pt")27        if os.path.exists(self.clip_model_path):28            print(f"Loading CLIP model from {self.clip_model_path}")29            self.clip_model, _ = clip.load(self.clip_model_path, device=self.device)30        else:31            print(f"CLIP model not found at {self.clip_model_path}, downloading...")32            self.clip_model, _ = clip.load("ViT-B-32", device=self.device)33        34        # Set model to evaluation mode35        self.clip_model.eval()36        37        print(f"DiffSketcher model initialized on device: {self.device}")38    39    def generate_svg(self, prompt, num_paths=10, width=512, height=512):40        """Generate an SVG from a text prompt"""41        print(f"Generating SVG for prompt: {prompt}")42        43        # Encode the prompt with CLIP44        with torch.no_grad():45            text_features = self.clip_model.encode_text(clip.tokenize([prompt]).to(self.device))46            text_features = text_features / text_features.norm(dim=-1, keepdim=True)47        48        # Generate a simple SVG based on the prompt49        # In a real implementation, this would use the full DiffSketcher model50        svg_content = f"""<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">51            <rect width="100%" height="100%" fill="#f0f0f0"/>52            <text x="50%" y="10%" font-family="Arial" font-size="20" text-anchor="middle">Generated by DiffSketcher</text>53            <text x="50%" y="50%" font-family="Arial" font-size="24" text-anchor="middle" font-weight="bold">{prompt}</text>54        """55        56        # Add some random paths based on the text features57        for i in range(min(num_paths, text_features.shape[1])):58            # Use the text features to generate path parameters59            feature_val = text_features[0, i % text_features.shape[1]].item()60            x = (feature_val + 1) * width / 261            y = ((i / num_paths) * 0.8 + 0.1) * height62            radius = abs(feature_val) * 50 + 1063            hue = (feature_val + 1) * 18064            65            # Add a circle with color based on the feature66            svg_content += f"""<circle cx="{x}" cy="{y}" r="{radius}" fill="hsl({hue}, 70%, 60%)" opacity="0.7" />"""67        68        # Close the SVG69        svg_content += "</svg>"70        71        return svg_content72    73    def svg_to_png(self, svg_content):74        """Convert SVG content to PNG"""75        try:76            png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))77            return png_data78        except Exception as e:79            print(f"Error converting SVG to PNG: {e}")80            # Create a simple error image81            image = Image.new("RGB", (512, 512), color="#ff0000")82            from PIL import ImageDraw83            draw = ImageDraw.Draw(image)84            draw.text((256, 256), f"Error: {str(e)}", fill="white", anchor="mm")85            86            # Convert PIL Image to PNG data87            buffer = io.BytesIO()88            image.save(buffer, format="PNG")89            return buffer.getvalue()90    91    def __call__(self, prompt):92        """Generate an SVG from a text prompt and convert to PNG"""93        svg_content = self.generate_svg(prompt)94        png_data = self.svg_to_png(svg_content)95        96        # Create a PIL Image from the PNG data97        image = Image.open(io.BytesIO(png_data))98        99        # Create the response100        response = {101            "svg": svg_content,102            "svg_base64": base64.b64encode(svg_content.encode("utf-8")).decode("utf-8"),103            "png_base64": base64.b64encode(png_data).decode("utf-8"),104            "image": image105        }106        107        return response