import gradio as gr import torch import os from PIL import Image import cairosvg import io import tempfile import argparse import gc import yaml import glob from huggingface_hub import hf_hub_download import spaces GLOBAL_PIPE = None device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM def parse_args(): parser = argparse.ArgumentParser(description='SVG Generator Service') parser.add_argument('--listen', type=str, default='0.0.0.0', help='Listen address (default: 0.0.0.0)') parser.add_argument('--port', type=int, default=7860, help='Port number (default: 7860)') parser.add_argument('--share', action='store_true', help='Enable gradio share link') parser.add_argument('--debug', action='store_true', help='Enable debug mode') return parser.parse_args() def load_models(): """Load models""" global GLOBAL_PIPE if GLOBAL_PIPE is None: # model_path = "/mnt/data-1/svg/model_outputs/hw_reason_0816" model_path = "Caffin/SVGThinker-7B" tok = AutoTokenizer.from_pretrained(model_path) model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="flash_attention_2" ) GLOBAL_PIPE = pipeline("text-generation", model=model, tokenizer=tok) import re def svg2png(svg_str, max_size=448): img_data = cairosvg.svg2png(bytestring=svg_str, output_height=448, background_color="white") image = Image.open(io.BytesIO(img_data)) if image.mode == "RGBA": new_img = Image.new("RGB", image.size, (255, 255, 255)) new_img.paste(image, mask=image.split()[3]) else: new_img = image.convert("RGB") width, height = new_img.size if width > height: new_width = max_size new_height = int((max_size / width) * height) else: new_height = max_size new_width = int((max_size / height) * width) new_img = new_img.resize((new_width, new_height)) return new_img def extract_svg(text): pattern = re.compile(r'', re.DOTALL) match = pattern.search(text) if match: return match.group() else: return None def generate_svg(text_description): try: gc.collect() torch.cuda.empty_cache() if torch.cuda.is_available() else None gen_config = dict( do_sample=True, temperature=0.8, top_p=0.6, top_k=100, repetition_penalty=1.05, num_beams=3, num_return_sequences=1, early_stopping=False, max_new_tokens=4096, ) if torch.cuda.is_available(): torch.cuda.synchronize() prompt = f"Review the given information below and generate a svg according to it.\n{text_description}" system_prompt = "You are a helpful assistant." prompt = f"SYSTEM: {system_prompt}\nUSER: {prompt}\nASSISTANT:" svg = GLOBAL_PIPE(prompt, **gen_config)[0]['generated_text'] # print(f'Raw SVG: {svg}') svg_str = svg svg_str = extract_svg(svg_str) # print(f"Extracted SVG: {svg_str}") png_image = svg2png(svg_str) # png_image.save(f'./debug_output.png') return svg_str, png_image except Exception as e: print(f"Generation error: {e}") import traceback traceback.print_exc() return f"Error: {e}", None @spaces.GPU(duration=75) def gradio_text_to_svg(text_description, progress=gr.Progress(track_tqdm=True)): """Gradio interface function - text-to-svg""" if not text_description or text_description.strip() == "": return "Please enter a description", None svg_code, png_image = generate_svg(text_description) return svg_code, png_image def create_interface(): # Example texts example_texts = [ '''The SVG represents a heart outlined in red on a white background. The heart is symmetrical with smooth curves and a pointed bottom. The outline is thick and consistent, but the shape is hollow inside with no fill color. The design conveys a classic and simple love symbol. ''', '''The SVG shows a computer monitor with a dark green screen containing four white, jagged lines representing data streams flowing from left to right. The monitor stand is colored in light gray. The rounded corners and flat design highlight a futuristic, tech-inspired style. ''', '''The SVG is a minimalist calendar icon. It has a white square with a thin black border and two black tabs protruding from the top edge. Inside the square, a centered black circle indicates a specific date, and below the calendar, a bold checkmark leans diagonally, symbolizing a completed event. The style is clean and monochrome. ''', '''The SVG features a flat illustration of a yellow rubber duck facing right. The duck has a round body, a small triangular orange beak, and a black dot for the eye. It sits on a shallow blue oval shape, suggesting water. The background is white and the design is clean and minimalistic.''', '''The SVG displays a black musical note icon on a white background. It is a solid quarter note with a round note head, vertical stem on the right, and a rightward flag curving upwards. The design is clean and modern with no textures or gradients, focusing on the iconic shape. ''', '''The SVG presents a minimalist side view of a sailboat with a single sail. The boat hull is colored dark blue, and the sail is a lighter blue with a triangular shape filled with a subtle wave pattern. The boat is floating on a thin, curved green line symbolizing water. The overall design is simple and stylized.''', '''The SVG is a geometric sun icon with eight orange rays radiating from a central yellow circle. The rays are trapezoidal, evenly spaced, and the circle has a subtle radial gradient from bright yellow at the center to light orange at the edges. A thin dark orange outline surrounds the whole shape. The design is symmetric and stylized with a crisp, flat aesthetic. ''', '''The SVG displays a single red heart shape centered on a soft pink rounded square background. The heart is filled with a smooth gradient from dark red at the bottom to lighter red at the top. The square background has rounded corners and a subtle drop shadow, giving the icon depth and softness. ''', '''The SVG is a rectangular icon featuring a gray document outline on a light beige background. A folded corner is present at the bottom left. Inside the document, there are four horizontal black lines symbolizing text content. The design is clean, simple, and modern. ''', '''The SVG shows a round magnifying glass with a blue handle and a transparent lens. Inside the lens is a simple black eye illustration with a blue iris and a black pupil. The handle has a slight gradient from light to dark blue, and the background is transparent, combining elements for a visual metaphor of close observation. ''', '''The SVG features a round emblem divided into three colored wedges: blue at the top left, yellow at the top right, and red at the bottom. Each wedge is separated by thin white lines, and the outer edge is framed with a subtle gray border. The design is simple and vibrant. ''', '''The SVG features a round, friendly emoticon with a thin orange outline circle on a white background. Inside the circle, two small solid orange dots represent eyes near the upper center. Below the eyes, a gentle upward curved orange line forms a smiling mouth, creating a cheerful expression. The style is clean and minimalist, using only orange and white, suitable for casual digital icons. ''', '''The SVG shows a large green checkmark angled diagonally across a white canvas. The stroke is thick and consistent, providing a strong visual contrast and a simple confirmation symbol without additional elements. ''', '''The SVG depicts a green circular icon with a white magnifying glass symbol in the center. The magnifying glass has a round lens and a short handle angled toward the bottom right. The colors are flat and solid, with the green background providing strong contrast to the white symbol. ''', '''The SVG is a simple icon of an envelope with a red outline and a white interior. The flap is indicated by a single diagonal line crossing the envelope front. The overall shape is rectangular with sharp corners, and the design uses a limited color palette emphasizing clarity and email-related symbolism. ''', '''The SVG illustrates a colorful hot air balloon set on a light blue circular background. The balloon is composed of vertical stripes in red, yellow, and orange hues, with a small brown basket hanging beneath. The overall shape is symmetrical and clean, emphasizing bright vivid colors for a joyful impression. ''', '''The SVG is a simple black outline of an over-ear headset with a modern, thin and curved headband connecting two oval ear cups. The ear cups have slight padding lines inside, and the overall icon is centered on a white background with thick black outlines.''', '''The SVG is a circular badge icon with a bold, red outline and a white interior. Centered within the circle is a black exclamation mark with a diamond-shaped dot beneath it. The design uses strong color contrast to command attention, typically serving as a notification or alert symbol in digital designs. ''', ] with gr.Blocks(title="SVGThinker-7B Demo Page", theme=gr.themes.Soft()) as demo: gr.Markdown("# SVGThinker-7B Demo Page") gr.Markdown("Generate an SVG from text.") with gr.Tabs(): # Text-to-SVG tab with gr.TabItem("Text-to-SVG"): with gr.Row(): with gr.Column(): text_input = gr.Textbox( label="Description", placeholder="Enter SVG description, e.g.: a circle", lines=3 ) # Add example texts gr.Examples( examples=[[text] for text in example_texts], inputs=[text_input], label="Example Descriptions (click to use)", examples_per_page=9 ) text_generate_btn = gr.Button("Generate SVG", variant="primary") with gr.Column(): text_svg_output = gr.Textbox( label="Generated SVG Code", lines=10, max_lines=20, show_copy_button=True ) text_png_preview = gr.Image(label="SVG Preview (PNG)", type="pil") text_generate_btn.click( fn=gradio_text_to_svg, inputs=[text_input], outputs=[text_svg_output, text_png_preview], queue=True ) return demo if __name__ == "__main__": # os.environ["TOKENIZERS_PARALLELISM"] = "true" args = parse_args() print("Loading models...") load_models() print("Models loaded successfully!") demo = create_interface() demo.launch( server_name=args.listen, server_port=args.port, share=args.share, debug=args.debug )