import logging
import re
import cairosvg
import torch
from transformers import AutoModelForCausalLM
from lxml import etree
import kagglehub
from gen_image import ImageGenerator
from starvector.data.util import process_and_rasterize_svg
svg_constraints = kagglehub.package_import('metric/svg-constraints')
class DLModel:
def __init__(self, model_id="starvector/starvector-8b-im2svg", device="cuda"):
"""
Initialize the SVG generation pipeline using StarVector.
Args:
model_id (str): The model identifier for the StarVector model.
device (str): The device to run the model on, either "cuda" or "cpu".
"""
self.image_generator = ImageGenerator(model_id="stabilityai/stable-diffusion-2-1-base", 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
# Load StarVector model
self.device = device
self.starvector = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
trust_remote_code=True
)
self.processor = self.starvector.model.processor
self.starvector.to(device)
self.starvector.eval()
def predict(self, description):
"""
Generate an SVG from a text description.
Args:
description (str): The text description to generate an image from.
Returns:
str: The generated SVG content.
"""
try:
# Step 1: Generate image using diffusion model
images = self.image_generator.generate(description)
image = images[0]
# Save the generated image
image_path = "diff_image.png"
image.save(image_path)
logging.info(f"Intermediate image saved to {image_path}")
# Step 2: Convert image to SVG using StarVector
processed_image = self.processor(image, return_tensors="pt")['pixel_values'].to(self.device)
if not processed_image.shape[0] == 1:
processed_image = processed_image.squeeze(0)
batch = {"image": processed_image}
with torch.no_grad():
raw_svg = self.starvector.generate_im2svg(batch, max_length=4000)[0]
raw_svg, _ = process_and_rasterize_svg(raw_svg)
if 'viewBox' not in raw_svg:
raw_svg = raw_svg.replace('<svg', f'<svg viewBox="0 0 384 384"')
# Step 3: Enforce constraints
svg_content = self.enforce_constraints(raw_svg)
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
# Example usage
if __name__ == "__main__":
model = DLModel()
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}")