Update app.py

app.py

@@ -14,6 +14,8 @@ import time
 import threading
 import spaces
 
+from huggingface_hub import hf_hub_download, snapshot_download
+
 from decoder import SketchDecoder
 from transformers import AutoTokenizer, AutoProcessor
 from qwen_vl_utils import process_vision_info
@@ -44,6 +46,10 @@ SUPPORTED_FORMATS = ['.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif']
 TARGET_IMAGE_SIZE = 448
 BLACK_COLOR_TOKEN = 40012
 
+# Default Hugging Face model IDs
+DEFAULT_QWEN_MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"
+DEFAULT_OMNISVG_MODEL = "OmniSVG/OmniSVG1.1_8B"
+
 # Task configurations with defaults
 TASK_CONFIGS = {
     "text-to-svg-icon": {
@@ -74,7 +80,6 @@ CUSTOM_CSS = """
     margin: 0 auto !important;
     padding: 20px !important;
 }
-
 /* Header styling */
 .header-container {
     text-align: center;
@@ -84,19 +89,16 @@ CUSTOM_CSS = """
     border-radius: 16px;
     color: white;
 }
-
 .header-container h1 {
     margin: 0;
     font-size: 2.5em;
     font-weight: 700;
 }
-
 .header-container p {
     margin: 10px 0 0 0;
     opacity: 0.9;
     font-size: 1.1em;
 }
-
 /* Tips section */
 .tips-box {
     background: #f8f9fa;
@@ -105,14 +107,12 @@ CUSTOM_CSS = """
     margin-bottom: 20px;
     border: 1px solid #e0e0e0;
 }
-
 .tips-box h3 {
     margin-top: 0;
     color: #333;
     border-bottom: 2px solid #667eea;
     padding-bottom: 10px;
 }
-
 .tip-category {
     background: white;
     border-radius: 8px;
@@ -120,19 +120,16 @@ CUSTOM_CSS = """
     margin: 10px 0;
     border-left: 4px solid #667eea;
 }
-
 .tip-category h4 {
     margin: 0 0 10px 0;
     color: #667eea;
 }
-
 .tip-category code {
     background: #f0f0f0;
     padding: 2px 6px;
     border-radius: 4px;
     font-size: 0.9em;
 }
-
 .example-prompt {
     background: #e8f4fd;
     padding: 10px;
@@ -142,12 +139,10 @@ CUSTOM_CSS = """
     font-size: 0.95em;
     color: #333;
 }
-
 .red-tip {
     color: #dc3545;
     font-weight: 600;
 }
-
 .red-box {
     background: #fff5f5;
     border: 1px solid #ffcccc;
@@ -156,11 +151,9 @@ CUSTOM_CSS = """
     border-radius: 8px;
     margin: 10px 0;
 }
-
 .red-box strong {
     color: #dc3545;
 }
-
 .orange-box {
     background: #fff8e6;
     border: 1px solid #ffc107;
@@ -169,11 +162,9 @@ CUSTOM_CSS = """
     border-radius: 8px;
     margin: 10px 0;
 }
-
 .orange-box strong {
     color: #ff9800;
 }
-
 .green-box {
     background: #e8f5e9;
     border: 1px solid #81c784;
@@ -182,21 +173,17 @@ CUSTOM_CSS = """
     border-radius: 8px;
     margin: 10px 0;
 }
-
 .green-box strong {
     color: #4caf50;
 }
-
 /* Tab styling */
 .tabs {
     border-radius: 12px !important;
     overflow: hidden;
 }
-
 .tabitem {
     padding: 20px !important;
 }
-
 /* Button styling */
 .primary-btn {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
@@ -205,12 +192,10 @@ CUSTOM_CSS = """
     padding: 12px 24px !important;
     font-size: 1.1em !important;
 }
-
 .primary-btn:hover {
     transform: translateY(-2px);
     box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
 }
-
 /* Settings group */
 .settings-group {
     background: #f8f9fa;
@@ -218,14 +203,12 @@ CUSTOM_CSS = """
     padding: 15px;
     margin: 10px 0;
 }
-
 .advanced-settings {
     background: #f0f4f8;
     border-radius: 8px;
     padding: 12px;
     margin-top: 10px;
 }
-
 /* Code output */
 .code-output textarea {
     font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace !important;
@@ -234,18 +217,15 @@ CUSTOM_CSS = """
     color: #d4d4d4 !important;
     border-radius: 8px !important;
 }
-
 /* Input image area */
 .input-image {
     border: 2px dashed #ccc;
     border-radius: 12px;
     transition: border-color 0.3s;
 }
-
 .input-image:hover {
     border-color: #667eea;
 }
-
 /* Footer */
 .footer {
     text-align: center;
@@ -253,7 +233,6 @@ CUSTOM_CSS = """
     color: #666;
     font-size: 0.9em;
 }
-
 /* Responsive adjustments */
 @media (max-width: 768px) {
     .gradio-container {
@@ -530,22 +509,85 @@ 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('--weight_path', type=str, default="/mnt/jfs-test/OmniSVG_result/8B_1126/1688_bs_4/merge_slerp/merge_150_350_bf16")
-    parser.add_argument('--model_path', type=str, default="/mnt/jfs-test/Qwen2.5-VL-7B-Instruct")
+    parser.add_argument('--weight_path', type=str, default=DEFAULT_OMNISVG_MODEL,
+                        help=f'HuggingFace repo ID or local path for OmniSVG weights (default: {DEFAULT_OMNISVG_MODEL})')
+    parser.add_argument('--model_path', type=str, default=DEFAULT_QWEN_MODEL,
+                        help=f'HuggingFace repo ID or local path for Qwen model (default: {DEFAULT_QWEN_MODEL})')
     return parser.parse_args()
 
 
-def load_models(weight_path, model_path):
-    """Load all models"""
+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:
+        local_path = hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            resume_download=True,
+        )
+        print(f"Successfully downloaded to: {local_path}")
+        return local_path
+    except Exception as e:
+        print(f"Error downloading from {repo_id}: {e}")
+        raise
+
+
+def is_local_path(path: str) -> bool:
+    """Check if a path is a local filesystem path or a HuggingFace repo ID."""
+    # If it contains a slash and is not a valid local path, assume it's a HF repo
+    if os.path.exists(path):
+        return True
+    # Check for common local path patterns
+    if path.startswith('/') or path.startswith('./') or path.startswith('../'):
+        return True
+    if os.path.sep in path and os.path.exists(os.path.dirname(path)):
+        return True
+    # Check Windows-style paths
+    if len(path) > 1 and path[1] == ':':
+        return True
+    return False
+
+
+def load_models(weight_path: str, model_path: str):
+    """
+    Load all models with support for both local paths and HuggingFace Hub.
+    
+    Args:
+        weight_path: Local path or HuggingFace repo ID for OmniSVG weights
+        model_path: Local path or HuggingFace repo ID for Qwen model
+    """
     global tokenizer, processor, sketch_decoder, svg_tokenizer
     
-    print(f"Loading models from {model_path}...")
+    print(f"Loading Qwen model from: {model_path}")
+    print(f"Loading OmniSVG weights from: {weight_path}")
     print(f"Using precision: {DTYPE}")
     
-    tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left")
-    processor = AutoProcessor.from_pretrained(model_path, padding_side="left")
+    # Load Qwen tokenizer and processor (transformers handles HF download automatically)
+    print("\n[1/3] Loading tokenizer and processor...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, 
+        padding_side="left",
+        trust_remote_code=True
+    )
+    processor = AutoProcessor.from_pretrained(
+        model_path, 
+        padding_side="left",
+        trust_remote_code=True
+    )
     processor.tokenizer.padding_side = "left"
+    print("Tokenizer and processor loaded successfully!")
 
+    # Initialize sketch decoder
+    print("\n[2/3] Initializing SketchDecoder...")
     sketch_decoder = SketchDecoder(
         pix_len=config['model']['max_length'],
         text_len=200,
@@ -553,17 +595,42 @@ def load_models(weight_path, model_path):
         torch_dtype=DTYPE
     )
     
-    bin_path = os.path.join(weight_path, "pytorch_model.bin")
-    if os.path.exists(bin_path):
-        print(f"Loading weights from: {bin_path}")
-        sketch_decoder.load_state_dict(torch.load(bin_path, map_location='cpu'))
+    # Load OmniSVG weights
+    print("\n[3/3] Loading OmniSVG weights...")
+    
+    # Check if weight_path is local or needs to be downloaded from HF
+    if is_local_path(weight_path):
+        # Local path - check for pytorch_model.bin
+        bin_path = os.path.join(weight_path, "pytorch_model.bin")
+        if not os.path.exists(bin_path):
+            # Maybe it's a direct path to the bin file
+            if os.path.exists(weight_path) and weight_path.endswith('.bin'):
+                bin_path = weight_path
+            else:
+                raise FileNotFoundError(
+                    f"Could not find pytorch_model.bin at {weight_path}. "
+                    f"Please provide a valid local path or HuggingFace repo ID."
+                )
+        print(f"Loading weights from local path: {bin_path}")
     else:
-        raise FileNotFoundError(f"No weights found at {bin_path}")
+        # HuggingFace repo ID - download the weights
+        print(f"Downloading weights from HuggingFace: {weight_path}")
+        bin_path = download_model_weights(weight_path, "pytorch_model.bin")
     
+    # Load the weights
+    state_dict = torch.load(bin_path, map_location='cpu')
+    sketch_decoder.load_state_dict(state_dict)
+    print("OmniSVG weights loaded successfully!")
+    
+    # Move to device and set eval mode
     sketch_decoder = sketch_decoder.to(device).eval()
+    
+    # Initialize SVG tokenizer
     svg_tokenizer = SVGTokenizer('./config.yaml')
     
+    print("\n" + "="*60)
     print("All models loaded successfully!")
+    print("="*60 + "\n")
 
 
 def detect_text_subtype(text_prompt):
@@ -731,7 +798,6 @@ def prepare_inputs(task_type, content):
         prompt_text = str(content).strip()
         
         instruction = f"""Generate an SVG illustration for: {prompt_text}
-
 Requirements:
 - Create complete SVG path commands
 - Include proper coordinates and colors
@@ -1354,12 +1420,13 @@ if __name__ == "__main__":
     print("="*60)
     print("OmniSVG Generator - Gradio App")
     print("="*60)
-    print(f"Model path: {args.model_path}")
-    print(f"Weight path: {args.weight_path}")
+    print(f"Qwen model: {args.model_path}")
+    print(f"OmniSVG weights: {args.weight_path}")
     print(f"Device: {device}")
+    print(f"Precision: {DTYPE}")
     print("="*60)
     
-    print("\nLoading models...")
+    print("\nLoading models (may download from HuggingFace Hub if needed)...")
     load_models(args.weight_path, args.model_path)
     print("Models loaded successfully!\n")