Update app.py

app.py

@@ -374,6 +374,7 @@ TIPS_HTML = """
         <ul style="margin: 8px 0 0 0; padding-left: 20px;">
             <li><strong>8B Model:</strong> Higher quality, more details, better for complex illustrations. Requires more VRAM (~16GB+).</li>
             <li><strong>4B Model:</strong> Faster, less VRAM required (~8GB+). Good for simple icons and basic shapes.</li>
+            <li><strong>Note:</strong> First generation with a new model size may take longer to load.</li>
         </ul>
     </div>
     
@@ -495,46 +496,6 @@ TIPS_HTML = """
         
     </div>
     
-    <!-- Extended Examples Section -->
-    <div style="margin-top: 20px; padding: 15px; background: #f0f7ff; border-radius: 10px; border: 1px solid #cce5ff;">
-        <h4 style="margin-top: 0; color: #0066cc;">More Complex Examples (Generate 6-8 candidates!)</h4>
-        
-        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)); gap: 12px; margin-top: 15px;">
-            <div class="example-prompt">
-                <strong>Business Avatar:</strong><br/>
-                "Circular professional avatar: man with short black hair, neutral skin tone round face, wearing dark navy suit with white shirt collar visible. Clean minimal style, centered in circle."
-            </div>
-            <div class="example-prompt">
-                <strong>Female Portrait:</strong><br/>
-                "Simple female face: oval face shape, long brown wavy hair on sides, two dot eyes, small nose, curved smile lips. Pink blush on cheeks. Cartoon portrait style."
-            </div>
-            <div class="example-prompt">
-                <strong>Child Character:</strong><br/>
-                "Cute child standing: large round head with short brown hair, big circular eyes with white highlights, small body in red t-shirt and blue shorts, simple stick arms and legs. Cheerful cartoon style."
-            </div>
-            <div class="example-prompt">
-                <strong>Active Pose:</strong><br/>
-                "Person walking: side view, circular head, rectangular torso in green jacket, legs in walking position (one forward, one back). Simple geometric style, moving right."
-            </div>
-            <div class="example-prompt">
-                <strong>Forest Scene:</strong><br/>
-                "Simple forest: light blue sky, row of 5 dark green triangular pine trees of varying heights, brown rectangular trunks, light green grass strip at bottom. Layered flat design."
-            </div>
-            <div class="example-prompt">
-                <strong>Ocean View:</strong><br/>
-                "Minimalist ocean: gradient blue sky at top, three horizontal wavy lines in dark blue for ocean, small white sailboat with triangular sail in center. Clean vector style."
-            </div>
-            <div class="example-prompt">
-                <strong>City Skyline:</strong><br/>
-                "Simple city skyline: orange sunset sky gradient, row of black rectangular building silhouettes of different heights, some with small yellow square windows. Minimalist style."
-            </div>
-            <div class="example-prompt">
-                <strong>Dog Character:</strong><br/>
-                "Friendly cartoon dog: brown oval body, round head with floppy ears, black dot nose, curved tail pointing up, four short legs. Sitting pose facing forward."
-            </div>
-        </div>
-    </div>
-    
     <!-- Quick Troubleshooting -->
     <div class="green-box" style="margin-top: 15px;">
         <strong>Quick Troubleshooting</strong>
@@ -582,26 +543,12 @@ def parse_args():
     parser.add_argument('--port', type=int, default=7860)
     parser.add_argument('--share', action='store_true')
     parser.add_argument('--debug', action='store_true')
-    parser.add_argument('--model_size', type=str, default=None,
-                        choices=AVAILABLE_MODEL_SIZES,
-                        help=f'Model size to load at startup (default: {DEFAULT_MODEL_SIZE}). Can be changed in UI.')
-    parser.add_argument('--weight_path', type=str, default=None,
-                        help='HuggingFace repo ID or local path for OmniSVG weights (overrides config)')
-    parser.add_argument('--model_path', type=str, default=None,
-                        help='HuggingFace repo ID or local path for Qwen model (overrides config)')
     return parser.parse_args()
 
 
 def download_model_weights(repo_id: str, filename: str = "pytorch_model.bin") -> str:
     """
     Download model weights from Hugging Face Hub.
-    
-    Args:
-        repo_id: Hugging Face repository ID (e.g., 'OmniSVG/OmniSVG1.1_8B')
-        filename: Name of the weights file to download
-        
-    Returns:
-        Local path to the downloaded file
     """
     print(f"Downloading {filename} from {repo_id}...")
     try:
@@ -633,11 +580,6 @@ def is_local_path(path: str) -> bool:
 def load_models(model_size: str, weight_path: str = None, model_path: str = None):
     """
     Load all models for a specific model size.
-    
-    Args:
-        model_size: Model size ("8B" or "4B")
-        weight_path: Local path or HuggingFace repo ID for OmniSVG weights (optional, uses config if None)
-        model_path: Local path or HuggingFace repo ID for Qwen model (optional, uses config if None)
     """
     global tokenizer, processor, sketch_decoder, svg_tokenizer, current_model_size
     
@@ -714,34 +656,34 @@ def load_models(model_size: str, weight_path: str = None, model_path: str = None
     print("="*60 + "\n")
 
 
-def switch_model(new_model_size: str):
+def ensure_model_loaded(model_size: str):
     """
-    Switch to a different model size.
-    
-    Args:
-        new_model_size: Target model size ("8B" or "4B")
-    
-    Returns:
-        Status message
+    Ensure the specified model is loaded. Load or switch if necessary.
+    This function should be called within @spaces.GPU decorated functions.
     """
-    global current_model_size
+    global current_model_size, sketch_decoder, tokenizer, processor, svg_tokenizer
     
-    if new_model_size == current_model_size:
-        return f"✅ Already using {new_model_size} model"
+    if current_model_size == model_size and sketch_decoder is not None:
+        return  # Already loaded
     
     with model_loading_lock:
-        # Clear memory
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        # Double-check after acquiring lock
+        if current_model_size == model_size and sketch_decoder is not None:
+            return
+        
+        # Clear old models if switching
+        if current_model_size is not None:
+            print(f"Switching from {current_model_size} to {model_size}...")
+            del sketch_decoder
+            del tokenizer
+            del processor
+            del svg_tokenizer
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
         
-        try:
-            load_models(new_model_size)
-            return f"✅ Successfully switched to {new_model_size} model"
-        except Exception as e:
-            error_msg = f"❌ Failed to switch to {new_model_size}: {str(e)}"
-            print(error_msg)
-            return error_msg
+        # Load new model
+        load_models(model_size)
 
 
 def detect_text_subtype(text_prompt):
@@ -1116,20 +1058,21 @@ def generate_candidates(inputs, task_type, subtype, temperature, top_p, top_k, r
     return all_candidates
 
 
-@spaces.GPU
-def gradio_text_to_svg(text_description, num_candidates, temperature, top_p, top_k, repetition_penalty, 
+@spaces.GPU(duration=300)  # 增加时间以支持模型加载
+def gradio_text_to_svg(text_description, model_size, num_candidates, temperature, top_p, top_k, repetition_penalty, 
                        progress=gr.Progress()):
-    """Gradio interface - text-to-svg with advanced parameters"""
+    """Gradio interface - text-to-svg with model selection"""
     if not text_description or text_description.strip() == "":
-        return '<div style="text-align:center;color:#999;padding:50px;">Please enter a description</div>', ""
+        return '<div style="text-align:center;color:#999;padding:50px;">Please enter a description</div>', "", f"Ready (no model loaded)"
     
     print("\n" + "="*60)
-    print(f"[TASK] text-to-svg (Model: {current_model_size})")
+    print(f"[TASK] text-to-svg")
+    print(f"[MODEL] Requested: {model_size}")
     print(f"[INPUT] {text_description[:100]}{'...' if len(text_description) > 100 else ''}")
     print(f"[PARAMS] candidates={num_candidates}, temp={temperature}, top_p={top_p}, top_k={top_k}, rep_penalty={repetition_penalty}")
     print("="*60)
     
-    progress(0, "Starting generation...")
+    progress(0, "Initializing...")
     
     gc.collect()
     if torch.cuda.is_available():
@@ -1137,9 +1080,16 @@ def gradio_text_to_svg(text_description, num_candidates, temperature, top_p, top
     
     start_time = time.time()
     
+    # 确保模型已加载
+    progress(0.02, f"Loading {model_size} model (first time may take a while)...")
+    ensure_model_loaded(model_size)
+    model_status = f"✅ Using {model_size} model"
+    
+    progress(0.05, "Model ready, starting generation...")
+    
     subtype = detect_text_subtype(text_description)
     print(f"[SUBTYPE] Detected: {subtype}")
-    progress(0.05, f"Detected: {subtype}")
+    progress(0.08, f"Detected: {subtype}")
     
     inputs = prepare_inputs("text-to-svg", text_description.strip())
     
@@ -1161,7 +1111,8 @@ def gradio_text_to_svg(text_description, num_candidates, temperature, top_p, top
         print("[WARNING] No valid SVG generated")
         return (
             '<div style="text-align:center;color:#999;padding:50px;">No valid SVG generated. Try different parameters or rephrase your prompt.</div>',
-            f"<!-- No valid SVG (took {elapsed:.1f}s) -->"
+            f"<!-- No valid SVG (took {elapsed:.1f}s) -->",
+            model_status
         )
     
     svg_codes = []
@@ -1174,28 +1125,30 @@ def gradio_text_to_svg(text_description, num_candidates, temperature, top_p, top
     progress(1.0, f"Done! {len(all_candidates)} candidates in {elapsed:.1f}s")
     print(f"[COMPLETE] text-to-svg finished\n")
     
-    return gallery_html, combined_svg
+    return gallery_html, combined_svg, model_status
 
 
-@spaces.GPU
-def gradio_image_to_svg(image, num_candidates, temperature, top_p, top_k, repetition_penalty,
+@spaces.GPU(duration=300)  # 增加时间以支持模型加载
+def gradio_image_to_svg(image, model_size, num_candidates, temperature, top_p, top_k, repetition_penalty,
                         replace_background, progress=gr.Progress()):
-    """Gradio interface - image-to-svg with background handling"""
+    """Gradio interface - image-to-svg with model selection"""
     
     if image is None:
         return (
             '<div style="text-align:center;color:#999;padding:50px;">Please upload an image</div>',
             "",
-            None
+            None,
+            f"Ready (no model loaded)"
         )
     
     print("\n" + "="*60)
-    print(f"[TASK] image-to-svg (Model: {current_model_size})")
+    print(f"[TASK] image-to-svg")
+    print(f"[MODEL] Requested: {model_size}")
     print(f"[INPUT] Image size: {image.size if hasattr(image, 'size') else 'unknown'}, mode: {image.mode if hasattr(image, 'mode') else 'unknown'}")
     print(f"[PARAMS] candidates={num_candidates}, temp={temperature}, top_p={top_p}, top_k={top_k}, rep_penalty={repetition_penalty}, replace_bg={replace_background}")
     print("="*60)
     
-    progress(0, "Processing input image...")
+    progress(0, "Initializing...")
     
     gc.collect()
     if torch.cuda.is_available():
@@ -1203,6 +1156,13 @@ def gradio_image_to_svg(image, num_candidates, temperature, top_p, top_k, repeti
     
     start_time = time.time()
     
+    # 确保模型已加载
+    progress(0.02, f"Loading {model_size} model (first time may take a while)...")
+    ensure_model_loaded(model_size)
+    model_status = f"✅ Using {model_size} model"
+    
+    progress(0.05, "Processing input image...")
+    
     img_processed, was_modified = preprocess_image_for_svg(
         image, 
         replace_background=replace_background,
@@ -1211,7 +1171,7 @@ def gradio_image_to_svg(image, num_candidates, temperature, top_p, top_k, repeti
     
     if was_modified:
         print("[PREPROCESS] Background processed/replaced")
-        progress(0.05, "Background processed")
+        progress(0.08, "Background processed")
     
     with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_file:
         img_processed.save(tmp_file.name, format='PNG', quality=100)
@@ -1240,7 +1200,8 @@ def gradio_image_to_svg(image, num_candidates, temperature, top_p, top_k, repeti
             return (
                 '<div style="text-align:center;color:#999;padding:50px;">No valid SVG generated. Try adjusting parameters.</div>',
                 f"<!-- No valid SVG (took {elapsed:.1f}s) -->",
-                img_processed
+                img_processed,
+                model_status
             )
         
         svg_codes = []
@@ -1253,7 +1214,7 @@ def gradio_image_to_svg(image, num_candidates, temperature, top_p, top_k, repeti
         progress(1.0, f"Done! {len(all_candidates)} candidates in {elapsed:.1f}s")
         print(f"[COMPLETE] image-to-svg finished\n")
         
-        return gallery_html, combined_svg, img_processed
+        return gallery_html, combined_svg, img_processed, model_status
         
     finally:
         if os.path.exists(tmp_path):
@@ -1323,7 +1284,7 @@ def create_interface():
     
     example_images = get_example_images()
     
-    with gr.Blocks(title="OmniSVG Demo Page") as demo:
+    with gr.Blocks(title="OmniSVG Generator", css=CUSTOM_CSS) as demo:
         # Header
         gr.HTML("""
         <div class="header-container">
@@ -1332,48 +1293,11 @@ def create_interface():
         </div>
         """)
         
-        # Model Selection Section
-        with gr.Row():
-            with gr.Column():
-                gr.HTML("""
-                <div class="blue-box">
-                    <strong>🔧 Model Selection</strong>
-                    <p style="margin: 5px 0 0 0; font-size: 0.9em;">
-                        Choose between <b>8B</b> (higher quality, more VRAM) or <b>4B</b> (faster, less VRAM).
-                    </p>
-                </div>
-                """)
-        
-        with gr.Row():
-            model_selector = gr.Dropdown(
-                choices=AVAILABLE_MODEL_SIZES,
-                value=DEFAULT_MODEL_SIZE,
-                label="Model Size",
-                info="8B: ~16GB VRAM, higher quality | 4B: ~8GB VRAM, faster",
-                interactive=True,
-                scale=1
-            )
-            model_status = gr.Textbox(
-                label="Model Status",
-                value=f"✅ Ready: {DEFAULT_MODEL_SIZE} model loaded",
-                interactive=False,
-                scale=2
-            )
-            switch_btn = gr.Button("Switch Model", variant="secondary", scale=1)
-        
-        # Model switch handler
-        switch_btn.click(
-            fn=switch_model,
-            inputs=[model_selector],
-            outputs=[model_status],
-            queue=True
-        )
-        
         # Queue status
         gr.HTML("""
         <div style="background: #e7f3ff; border: 1px solid #b3d7ff; border-radius: 8px; padding: 12px 15px; margin: 15px 0;">
             <span style="font-size: 1.5em;">ℹ️</span>
-            <strong>Queue System Active</strong> - Requests processed one at a time. Please wait patiently if busy.
+            <strong>Queue System Active</strong> - Requests processed one at a time. First generation with a new model may take longer to load.
         </div>
         """)
         
@@ -1399,6 +1323,15 @@ def create_interface():
                         
                         with gr.Group(elem_classes=["settings-group"]):
                             gr.Markdown("### Settings")
+                            
+                            # Model Selection
+                            img_model_size = gr.Dropdown(
+                                choices=AVAILABLE_MODEL_SIZES,
+                                value=DEFAULT_MODEL_SIZE,
+                                label="Model Size",
+                                info="8B: Higher quality (~16GB VRAM) | 4B: Faster (~8GB VRAM)"
+                            )
+                            
                             img_num_candidates = gr.Slider(
                                 minimum=1, maximum=MAX_NUM_CANDIDATES, value=DEFAULT_NUM_CANDIDATES, step=1,
                                 label="Number of Candidates"
@@ -1443,6 +1376,13 @@ def create_interface():
                             elem_classes=["primary-btn"]
                         )
                         
+                        # Model status display
+                        img_model_status = gr.Textbox(
+                            label="Model Status",
+                            value="Ready (model loads on first generation)",
+                            interactive=False
+                        )
+                        
                         if example_images:
                             gr.Markdown("### Examples")
                             gr.Examples(examples=example_images, inputs=[image_input], label="")
@@ -1461,9 +1401,9 @@ def create_interface():
                 
                 image_generate_btn.click(
                     fn=gradio_image_to_svg,
-                    inputs=[image_input, img_num_candidates, img_temperature, img_top_p, 
+                    inputs=[image_input, img_model_size, img_num_candidates, img_temperature, img_top_p, 
                            img_top_k, img_rep_penalty, img_replace_bg],
-                    outputs=[image_gallery, image_svg_output, image_processed],
+                    outputs=[image_gallery, image_svg_output, image_processed, img_model_status],
                     queue=True
                 )
             
@@ -1485,6 +1425,15 @@ def create_interface():
                         
                         with gr.Group(elem_classes=["settings-group"]):
                             gr.Markdown("### Settings")
+                            
+                            # Model Selection
+                            text_model_size = gr.Dropdown(
+                                choices=AVAILABLE_MODEL_SIZES,
+                                value=DEFAULT_MODEL_SIZE,
+                                label="Model Size",
+                                info="8B: Higher quality (~16GB VRAM) | 4B: Faster (~8GB VRAM)"
+                            )
+                            
                             text_num_candidates = gr.Slider(
                                 minimum=1, maximum=MAX_NUM_CANDIDATES, value=6, step=1,
                                 label="Number of Candidates",
@@ -1526,6 +1475,13 @@ def create_interface():
                             elem_classes=["primary-btn"]
                         )
                         
+                        # Model status display
+                        text_model_status = gr.Textbox(
+                            label="Model Status",
+                            value="Ready (model loads on first generation)",
+                            interactive=False
+                        )
+                        
                         gr.Markdown("### Example Prompts (30)")
                         gr.Examples(
                             examples=[[text] for text in example_texts],
@@ -1549,16 +1505,16 @@ def create_interface():
                 
                 text_generate_btn.click(
                     fn=gradio_text_to_svg,
-                    inputs=[text_input, text_num_candidates, text_temperature, text_top_p, 
+                    inputs=[text_input, text_model_size, text_num_candidates, text_temperature, text_top_p, 
                            text_top_k, text_rep_penalty],
-                    outputs=[text_gallery, text_svg_output],
+                    outputs=[text_gallery, text_svg_output, text_model_status],
                     queue=True
                 )
         
         # Footer
         gr.HTML(f"""
         <div class="footer">
-            <p>Built with OmniSVG | Current Model: <strong>{DEFAULT_MODEL_SIZE}</strong></p>
+            <p>Built with OmniSVG | Available models: {', '.join(AVAILABLE_MODEL_SIZES)}</p>
             <p style="color: #dc3545; font-weight: 600;">Remember: Generate 4-8 candidates and pick the best!</p>
         </div>
         """)
@@ -1571,16 +1527,14 @@ if __name__ == "__main__":
     
     args = parse_args()
     
-    # Determine initial model size
-    initial_model_size = args.model_size or DEFAULT_MODEL_SIZE
-    
     print("="*60)
-    print("OmniSVG Demo Page - Gradio App")
+    print("OmniSVG Demo Page - Gradio App (HuggingFace Space)")
     print("="*60)
     print(f"Available model sizes: {AVAILABLE_MODEL_SIZES}")
-    print(f"Initial model size: {initial_model_size}")
+    print(f"Default model size: {DEFAULT_MODEL_SIZE}")
     print(f"Device: {device}")
     print(f"Precision: {DTYPE}")
+    print("Models will be loaded on-demand when first used.")
     print("="*60)
     
     # Print loaded config values
@@ -1594,10 +1548,6 @@ if __name__ == "__main__":
     print(f"  - PAD_TOKEN_ID: {PAD_TOKEN_ID}")
     print("="*60)
     
-    print("\nLoading models (may download from HuggingFace Hub if needed)...")
-    load_models(initial_model_size, args.weight_path, args.model_path)
-    print("Models loaded successfully!\n")
-    
     demo = create_interface()
     
     demo.queue(default_concurrency_limit=1, max_size=20)