Update app.py

app.py

@@ -1,141 +1,61 @@
-import argparse
 import gradio as gr
-import os
-import torch
-import trimesh
-import sys
-from pathlib import Path
-
-pathdir = Path(__file__).parent / 'cube'
-sys.path.append(pathdir.as_posix())
-
-# print(__file__)
-# print(os.listdir())
-# print(os.listdir('cube'))
-# print(pathdir.as_posix())
-
-from cube3d.inference.engine import EngineFast, Engine
-from pathlib import Path
-import uuid
-import shutil
-from huggingface_hub import snapshot_download
-
-
-GLOBAL_STATE = {}
-
-def gen_save_folder(max_size=200):
-    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
-
-    dirs = [f for f in Path(GLOBAL_STATE["SAVE_DIR"]).iterdir() if f.is_dir()]
-
-    if len(dirs) >= max_size:
-        oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime)
-        shutil.rmtree(oldest_dir)
-        print(f"Removed the oldest folder: {oldest_dir}")
-
-    new_folder = os.path.join(GLOBAL_STATE["SAVE_DIR"], str(uuid.uuid4()))
-    os.makedirs(new_folder, exist_ok=True)
-    print(f"Created new folder: {new_folder}")
-
-    return new_folder
-
-def handle_text_prompt(input_prompt, variance = 0):
-    print(f"prompt: {input_prompt}, variance: {variance}")
-    top_p = None if variance == 0 else (100 - variance) / 100.0
-    mesh_v_f = GLOBAL_STATE["engine_fast"].t2s([input_prompt], use_kv_cache=True, resolution_base=8.0, top_p=top_p)
-    # save output
-    vertices, faces = mesh_v_f[0][0], mesh_v_f[0][1]
-    save_folder = gen_save_folder()
-    output_path = os.path.join(save_folder, "output.glb")
-    trimesh.Trimesh(vertices=vertices, faces=faces).export(output_path)
-    return output_path
-
-def build_interface():
-    """Build UI for gradio app
-    """
-    title = "Cube 3D"
-    with gr.Blocks(theme=gr.themes.Soft(), title=title, fill_width=True) as interface:
-        gr.Markdown(
-            f"""
-            # {title}
-            # Check out our [Github](https://github.com/Roblox/cube) to try it on your own machine!
-            """
-        )
-
-        with gr.Row():
-            with gr.Column(scale=2):
-                with gr.Group():
-                    input_text_box = gr.Textbox(
-                        value=None,
-                        label="Prompt",
-                        lines=2,
-                    )
-                    variance = gr.Slider(minimum=0, maximum=99, step=1, value=0, label="Variance")
-                with gr.Row():
-                    submit_button = gr.Button("Submit", variant="primary")
-            with gr.Column(scale=3):
-                model3d = gr.Model3D(
-                    label="Output", height="45em", interactive=False
-                )
-    
-        submit_button.click(
-            handle_text_prompt,
-            inputs=[
-                input_text_box,
-                variance
-            ],
-            outputs=[
-                model3d
-            ]
-        )
-                
-    return interface
-
-if __name__=="__main__":
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--config_path",
-        type=str,
-        help="Path to the config file",
-        default="cube/cube3d/configs/open_model.yaml",
-    )
-    parser.add_argument(
-        "--gpt_ckpt_path",
-        type=str,
-        help="Path to the gpt ckpt path",
-        default="model_weights/shape_gpt.safetensors",
-    )
-    parser.add_argument(
-        "--shape_ckpt_path",
-        type=str,
-        help="Path to the shape ckpt path",
-        default="model_weights/shape_tokenizer.safetensors",
-    )
-    parser.add_argument(
-        "--save_dir",
-        type=str,
-        default="gradio_save_dir",
-    )
-
-    args = parser.parse_args()
-    snapshot_download(
-        repo_id="Roblox/cube3d-v0.1",
-        local_dir="./model_weights"
-    )
-    config_path = args.config_path
-    gpt_ckpt_path = "./model_weights/shape_gpt.safetensors"
-    shape_ckpt_path = "./model_weights/shape_tokenizer.safetensors"
-    engine_fast = EngineFast(
-        config_path,
-        gpt_ckpt_path, 
-        shape_ckpt_path,
-        device=torch.device("cuda"),
-    )
-    GLOBAL_STATE["engine_fast"] = engine_fast
-    GLOBAL_STATE["SAVE_DIR"] = args.save_dir
-    os.makedirs(GLOBAL_STATE["SAVE_DIR"], exist_ok=True)
-
-    demo = build_interface()
-    demo.queue(default_concurrency_limit=1)
-    demo.launch()
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import svgwrite
+import cairosvg
+import speech_recognition as sr
+import io
+
+# Load the StarVector model
+tokenizer = AutoTokenizer.from_pretrained("starvector/starvector-8b-im2svg")
+model = AutoModelForCausalLM.from_pretrained("starvector/starvector-8b-im2svg")
+
+def generate_svg(prompt, width, height):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=512)
+    svg_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Ensure SVG is properly wrapped
+    svg_wrapped = f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">{svg_code}</svg>'
+
+    # Convert to PNG
+    png_output = cairosvg.svg2png(bytestring=svg_wrapped.encode('utf-8'))
+
+    with open("output.svg", "w") as f:
+        f.write(svg_wrapped)
+    with open("output.png", "wb") as f:
+        f.write(png_output)
+
+    return svg_wrapped, "output.png", "output.svg"
+
+def transcribe_audio(audio_path):
+    recognizer = sr.Recognizer()
+    with sr.AudioFile(audio_path) as source:
+        audio_data = recognizer.record(source)
+    return recognizer.recognize_google(audio_data)
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Vector Logo Generator (Text + Voice)")
+
+    with gr.Row():
+        txt = gr.Textbox(label="Text Prompt")
+        mic = gr.Audio(source="microphone", type="filepath", label="Or speak your prompt")
+
+    with gr.Row():
+        width = gr.Slider(minimum=100, maximum=1000, value=500, step=10, label="Width (px)")
+        height = gr.Slider(minimum=100, maximum=1000, value=500, step=10, label="Height (px)")
+
+    svg_output = gr.Textbox(label="SVG Code Output")
+    png_output = gr.Image(label="PNG Preview")
+    svg_file = gr.File(label="Download SVG")
+    png_file = gr.File(label="Download PNG")
+
+    def run(prompt, audio, w, h):
+        if not prompt and audio:
+            prompt = transcribe_audio(audio)
+        svg, png_path, svg_path = generate_svg(prompt, w, h)
+        return svg, png_path, svg_path
+
+    run_button = gr.Button("Generate")
+    run_button.click(fn=run, inputs=[txt, mic, width, height], outputs=[svg_output, png_output, svg_file, png_file])
+
+demo.launch()