Spaces:
Running
Running
File size: 9,785 Bytes
6642f4e da1ac10 6642f4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 |
import os
import tempfile
import logging
import re
import subprocess
import cairosvg
from lxml import etree
import kagglehub
from gen_image import ImageGenerator
import vtracer
svg_constraints = kagglehub.package_import('metric/svg-constraints')
class MLModel:
def __init__(self, model_id="stabilityai/stable-diffusion-2-1-base", device="cuda"):
"""
Initialize the SVG generation pipeline.
Args:
model_id (str): The model identifier for the stable diffusion model.
device (str): The device to run the model on, either "cuda" or "cpu".
"""
self.image_generator = ImageGenerator(model_id=model_id, device=device)
self.default_svg = """<svg width="256" height="256" viewBox="0 0 256 256"><circle cx="50" cy="50" r="40" fill="red" /></svg>"""
self.constraints = svg_constraints.SVGConstraints()
self.timeout_seconds = 90
def predict(self, description, simplify=True, color_precision=6,
filter_speckle=4, path_precision=8):
"""
Generate an SVG from a text description.
Args:
description (str): The text description to generate an image from.
simplify (bool): Whether to simplify the SVG paths.
color_precision (int): Color quantization precision.
filter_speckle (int): Filter speckle size.
path_precision (int): Path fitting precision.
Returns:
str: The generated SVG content.
"""
try:
# Step 1: Generate image using diffusion model
images = self.image_generator.generate(description)
image = images[0]
# Step 2: Save image to a temporary file
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_img:
temp_img_path = temp_img.name
image.save(temp_img_path)
# Step 3: Convert image to SVG using vtracer
with tempfile.NamedTemporaryFile(suffix='.svg', delete=False) as temp_svg:
temp_svg_path = temp_svg.name
# Process the image with vtracer using parameters directly
vtracer.convert_image_to_svg_py(
temp_img_path,
temp_svg_path,
colormode='color',
hierarchical='stacked' if simplify else 'cutout',
mode='spline',
filter_speckle=filter_speckle,
color_precision=color_precision,
path_precision=path_precision,
corner_threshold=60,
length_threshold=4.0,
max_iterations=10,
splice_threshold=45
)
# Step 4: Read the generated SVG
with open(temp_svg_path, 'r') as f:
svg_content = f.read()
# Clean up temporary files
os.unlink(temp_img_path)
os.unlink(temp_svg_path)
# Step 5: Enforce constraints
svg_content = self.enforce_constraints(svg_content)
return svg_content
except Exception as e:
logging.error(f"Error generating SVG: {e}")
return self.default_svg
def enforce_constraints(self, svg_string: str) -> str:
"""Enforces constraints on an SVG string, removing disallowed elements
and attributes.
Parameters
----------
svg_string : str
The SVG string to process.
Returns
-------
str
The processed SVG string, or the default SVG if constraints
cannot be satisfied.
"""
logging.info('Sanitizing SVG...')
try:
# Remove XML declaration if it exists
svg_string = re.sub(r'<\?xml[^>]+\?>', '', svg_string).strip()
parser = etree.XMLParser(remove_blank_text=True, remove_comments=True)
root = etree.fromstring(svg_string, parser=parser)
except etree.ParseError as e:
logging.error('SVG Parse Error: %s. Returning default SVG.', e)
logging.error('SVG string: %s', svg_string)
return self.default_svg
elements_to_remove = []
for element in root.iter():
tag_name = etree.QName(element.tag).localname
# Remove disallowed elements
if tag_name not in self.constraints.allowed_elements:
elements_to_remove.append(element)
continue # Skip attribute checks for removed elements
# Remove disallowed attributes
attrs_to_remove = []
for attr in element.attrib:
attr_name = etree.QName(attr).localname
if (
attr_name
not in self.constraints.allowed_elements[tag_name]
and attr_name
not in self.constraints.allowed_elements['common']
):
attrs_to_remove.append(attr)
for attr in attrs_to_remove:
logging.debug(
'Attribute "%s" for element "%s" not allowed. Removing.',
attr,
tag_name,
)
del element.attrib[attr]
# Check and remove invalid href attributes
for attr, value in element.attrib.items():
if etree.QName(attr).localname == 'href' and not value.startswith('#'):
logging.debug(
'Removing invalid href attribute in element "%s".', tag_name
)
del element.attrib[attr]
# Validate path elements to help ensure SVG conversion
if tag_name == 'path':
d_attribute = element.get('d')
if not d_attribute:
logging.warning('Path element is missing "d" attribute. Removing path.')
elements_to_remove.append(element)
continue # Skip further checks for this removed element
# Use regex to validate 'd' attribute format
path_regex = re.compile(
r'^' # Start of string
r'(?:' # Non-capturing group for each command + numbers block
r'[MmZzLlHhVvCcSsQqTtAa]' # Valid SVG path commands (adjusted to exclude extra letters)
r'\s*' # Optional whitespace after command
r'(?:' # Non-capturing group for optional numbers
r'-?\d+(?:\.\d+)?(?:[Ee][+-]?\d+)?' # First number
r'(?:[\s,]+-?\d+(?:\.\d+)?(?:[Ee][+-]?\d+)?)*' # Subsequent numbers with mandatory separator(s)
r')?' # Numbers are optional (e.g. for Z command)
r'\s*' # Optional whitespace after numbers/command block
r')+' # One or more command blocks
r'\s*' # Optional trailing whitespace
r'$' # End of string
)
if not path_regex.match(d_attribute):
logging.warning(
'Path element has malformed "d" attribute format. Removing path.'
)
elements_to_remove.append(element)
continue
logging.debug('Path element "d" attribute validated (regex check).')
# Remove elements marked for removal
for element in elements_to_remove:
if element.getparent() is not None:
element.getparent().remove(element)
logging.debug('Removed element: %s', element.tag)
try:
cleaned_svg_string = etree.tostring(root, encoding='unicode', xml_declaration=False)
return cleaned_svg_string
except ValueError as e:
logging.error(
'SVG could not be sanitized to meet constraints: %s', e
)
return self.default_svg
def optimize_svg(self, svg_content):
"""
Optimize the SVG content using SVGO.
Args:
svg_content (str): The SVG content to optimize.
Returns:
str: The optimized SVG content.
"""
try:
with tempfile.NamedTemporaryFile(suffix='.svg', delete=False) as temp_svg:
temp_svg_path = temp_svg.name
temp_svg.write(svg_content.encode('utf-8'))
with tempfile.NamedTemporaryFile(suffix='.svg', delete=False) as temp_out:
temp_out_path = temp_out.name
subprocess.run(["svgo", temp_svg_path, "-o", temp_out_path], check=True)
with open(temp_out_path, 'r') as f:
optimized_svg = f.read()
os.unlink(temp_svg_path)
os.unlink(temp_out_path)
return optimized_svg
except (FileNotFoundError, subprocess.CalledProcessError):
print("Warning: SVGO not found or failed. Returning unoptimized SVG.")
return svg_content
# Example usage
if __name__ == "__main__":
model = MLModel()
svg = model.predict("a purple forest at dusk")
# Convert SVG to PNG
try:
# Create a PNG in memory
png_data = cairosvg.svg2png(bytestring=svg.encode('utf-8'))
# Save the PNG to a file
with open("output.png", "wb") as f:
f.write(png_data)
print("SVG saved as output.png")
except Exception as e:
print(f"Error converting SVG to PNG: {e}") |