Update app.py

app.py

@@ -1,234 +1,112 @@
-import spaces
 import gradio as gr
-import sys
-import threading
-import queue
-from io import TextIOBase
-from inference import inference_patch
-import datetime
-import subprocess
-import os
-
-# Predefined valid combinations set
-with open('prompts.txt', 'r') as f:
-    prompts = f.readlines()
-valid_combinations = set()
-for prompt in prompts:
-    prompt = prompt.strip()
-    parts = prompt.split('_')
-    valid_combinations.add((parts[0], parts[1], parts[2]))
-
-# Generate available options
-periods = sorted({p for p, _, _ in valid_combinations})
-composers = sorted({c for _, c, _ in valid_combinations})
-instruments = sorted({i for _, _, i in valid_combinations})
-
-# Dynamic component updates
-def update_components(period, composer):
-    if not period:
-        return [
-            gr.Dropdown(choices=[], value=None, interactive=False),
-            gr.Dropdown(choices=[], value=None, interactive=False)
-        ]
-    
-    valid_composers = sorted({c for p, c, _ in valid_combinations if p == period})
-    valid_instruments = sorted({i for p, c, i in valid_combinations if p == period and c == composer}) if composer else []
-    
-    return [
-        gr.Dropdown(
-            choices=valid_composers,
-            value=composer if composer in valid_composers else None,
-            interactive=True  
-        ),
-        gr.Dropdown(
-            choices=valid_instruments,
-            value=None,
-            interactive=bool(valid_instruments)  
-        )
-    ]
-
-
-class RealtimeStream(TextIOBase):
-    def __init__(self, queue):
-        self.queue = queue
-    
-    def write(self, text):
-        self.queue.put(text)
-        return len(text)
-
-
-def save_and_convert(abc_content, period, composer, instrumentation):
-    if not all([period, composer, instrumentation]):
-        raise gr.Error("Please complete a valid generation first before saving")
-    
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    prompt_str = f"{period}_{composer}_{instrumentation}"
-    filename_base = f"{timestamp}_{prompt_str}"
-    
-    abc_filename = f"{filename_base}.abc"
-    with open(abc_filename, "w", encoding="utf-8") as f:
-        f.write(abc_content)
-
-    xml_filename = f"{filename_base}.xml"
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoProcessor
+from starvector.data.util import process_and_rasterize_svg
+import torch
+import io
+
+USE_BOTH_MODELS = True  # Set this to True to load both models
+
+# Load models at startup
+models = {}
+if USE_BOTH_MODELS:
+    # Load 8b model
+    model_name_8b = "starvector/starvector-8b-im2svg"
+    models['8b'] = {
+        'model': AutoModelForCausalLM.from_pretrained(model_name_8b, torch_dtype=torch.float16, trust_remote_code=True),
+        'processor': None  # Will be set below
+    }
+    models['8b']['model'].cuda()
+    models['8b']['model'].eval()
+    models['8b']['processor'] = models['8b']['model'].model.processor
+
+    # Load 1b model
+    model_name_1b = "starvector/starvector-1b-im2svg"
+    models['1b'] = {
+        'model': AutoModelForCausalLM.from_pretrained(model_name_1b, torch_dtype=torch.float16, trust_remote_code=True),
+        'processor': None
+    }
+    models['1b']['model'].cuda()
+    models['1b']['model'].eval()
+    models['1b']['processor'] = models['1b']['model'].model.processor
+else:
+    # Load only 8b model
+    model_name = "starvector/starvector-8b-im2svg"
+    models['8b'] = {
+        'model': AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True),
+        'processor': None
+    }
+    models['8b']['model'].cuda()
+    models['8b']['model'].eval()
+    models['8b']['processor'] = models['8b']['model'].model.processor
+
+def convert_to_svg(image, model_choice):
     try:
-        subprocess.run(
-            ["python", "abc2xml.py", '-o', '.', abc_filename, ],
-            check=True,
-            capture_output=True,
-            text=True
-        )
-    except subprocess.CalledProcessError as e:
-        error_msg = f"Conversion failed: {e.stderr}" if e.stderr else "Unknown error"
-        raise gr.Error(f"ABC to XML conversion failed: {error_msg}. Please try to generate another composition.")
-    
-    return f"Saved successfully: {abc_filename} -> {xml_filename}"
-
-
-@spaces.GPU
-def generate_music(period, composer, instrumentation):
-    if (period, composer, instrumentation) not in valid_combinations:
-        raise gr.Error("Invalid prompt combination! Please re-select from the period options")
-    
-    output_queue = queue.Queue()
-    original_stdout = sys.stdout
-    sys.stdout = RealtimeStream(output_queue)
-    
-    result_container = []
-    def run_inference():
-        try:
-            result_container.append(inference_patch(period, composer, instrumentation))
-        finally:
-            sys.stdout = original_stdout
-    
-    thread = threading.Thread(target=run_inference)
-    thread.start()
-    
-    process_output = ""
-    while thread.is_alive():
-        try:
-            text = output_queue.get(timeout=0.1)
-            process_output += text
-            yield process_output, None  
-        except queue.Empty:
-            continue
-    
-    while not output_queue.empty():
-        text = output_queue.get()
-        process_output += text
-        yield process_output, None
-    
-    final_result = result_container[0] if result_container else ""
-    yield process_output, final_result
-
-with gr.Blocks() as demo:
-    gr.Markdown("## NotaGen")
+        if image is None:
+            return None, None, "Please upload an image first"
+        
+        # Select the model based on user choice
+        selected_model = models[model_choice]['model']
+        selected_processor = models[model_choice]['processor']
+        
+        # Process the uploaded image
+        image_pil = Image.open(image)
+        image_tensor = selected_processor(image_pil, return_tensors="pt")['pixel_values'].cuda()
+        
+        if not image_tensor.shape[0] == 1:
+            image_tensor = image_tensor.squeeze(0)
+        
+        batch = {"image": image_tensor}
+        
+        # Generate SVG
+        raw_svg = selected_model.generate_im2svg(batch, max_length=4000)[0]
+        svg, raster_image = process_and_rasterize_svg(raw_svg)
+        
+        # Convert SVG string to bytes for download
+        svg_bytes = io.BytesIO(svg.encode('utf-8'))
+        
+        return raster_image, svg_bytes, f"Conversion successful using {model_choice} model!"
+    except Exception as e:
+        return None, None, f"Error: {str(e)}"
+
+# Create Blocks interface
+with gr.Blocks(title="Image to SVG Converter") as demo:
+    gr.Markdown("# Image to SVG Converter")
+    gr.Markdown("Upload an image to convert it to SVG format using StarVector model")
     
     with gr.Row():
-        # 左侧栏
-        with gr.Column():
-            period_dd = gr.Dropdown(
-                choices=periods,
-                value=None, 
-                label="Period",
-                interactive=True
-            )
-            composer_dd = gr.Dropdown(
-                choices=[],
-                value=None,
-                label="Composer",
-                interactive=False
-            )
-            instrument_dd = gr.Dropdown(
-                choices=[],
-                value=None,
-                label="Instrumentation",
-                interactive=False
-            )
-            
-            generate_btn = gr.Button("Generate!", variant="primary")
-            
-            process_output = gr.Textbox(
-                label="Generation process",
-                interactive=False,
-                lines=15,
-                max_lines=15,
-                placeholder="Generation progress will be shown here...",
-                elem_classes="process-output"
+        with gr.Column(scale=1):
+            # Input section
+            input_image = gr.Image(type="filepath", label="Upload Image")
+            if USE_BOTH_MODELS:
+                model_choice = gr.Radio(
+                    choices=["8b", "1b"],
+                    value="8b",
+                    label="Select Model",
+                    info="Choose between 8b (larger) and 1b (smaller) models"
+                )
+            convert_btn = gr.Button("Convert to SVG")
+            example = gr.Examples(
+                examples=[["assets/examples/sample-18.png"]],
+                inputs=input_image
             )
-
-        # 右侧栏
-        with gr.Column():
-            final_output = gr.Textbox(
-                label="Post-processed ABC notation scores",
-                interactive=True,
-                lines=23,
-                placeholder="Post-processed ABC scores will be shown here...",
-                elem_classes="final-output"
-            )
-            
-            with gr.Row():
-                save_btn = gr.Button("💾 Save as ABC & XML files", variant="secondary")
-            
-            save_status = gr.Textbox(
-                label="Save Status",
-                interactive=False,
-                visible=True,
-                max_lines=2
-            )
-    
-    period_dd.change(
-        update_components,
-        inputs=[period_dd, composer_dd],
-        outputs=[composer_dd, instrument_dd]
-    )
-    composer_dd.change(
-        update_components,
-        inputs=[period_dd, composer_dd],
-        outputs=[composer_dd, instrument_dd]
+        
+        with gr.Column(scale=1):
+            # Output section
+            output_preview = gr.Image(type="pil", label="Rasterized SVG Preview")
+            output_file = gr.File(label="Download SVG")
+            status = gr.Textbox(label="Status")
+    
+    # Connect button click to conversion function
+    inputs = [input_image]
+    if USE_BOTH_MODELS:
+        inputs.append(model_choice)
+        
+    convert_btn.click(
+        fn=convert_to_svg,
+        inputs=inputs,
+        outputs=[output_preview, output_file, status]
     )
-    
-    generate_btn.click(
-        generate_music,
-        inputs=[period_dd, composer_dd, instrument_dd],
-        outputs=[process_output, final_output]
-    )
-    
-    save_btn.click(
-        save_and_convert,
-        inputs=[final_output, period_dd, composer_dd, instrument_dd],
-        outputs=[save_status]
-    )
-
-
-css = """
-.process-output {
-    background-color: #f0f0f0;
-    font-family: monospace;
-    padding: 10px;
-    border-radius: 5px;
-}
-.final-output {
-    background-color: #ffffff;
-    font-family: sans-serif;
-    padding: 10px;
-    border-radius: 5px;
-}
-
-.process-output textarea {
-    max-height: 500px !important;
-    overflow-y: auto !important;
-    white-space: pre-wrap;
-}
-
-"""
-css += """
-button#💾-save-convert:hover {
-    background-color: #ffe6e6;
-}
-"""
-
-demo.css = css
-
-if __name__ == "__main__":
 
-    demo.queue().launch()
+# Launch the app
+demo.launch()