Ensure --model argument adds model to database

- Added ensure_model_exists function to database.py
- Modified vidai.py to call ensure_model_exists when --model is specified
- Automatically detects Hugging Face vs local models

vidai.py

@@ -44,6 +44,7 @@ from vidai.config import (
     get_backend_host, set_backend_host, get_backend_web_port, set_backend_web_port,
     get_backend_worker_port, set_backend_worker_port
 )
+from vidai.database import ensure_model_exists
 from vidai.cluster_master import start_cluster_master
 
 def main():
@@ -323,6 +324,11 @@ Examples:
     # Update config with command line values
     set_default_model(args.model)
     set_default_model_type(args.model_type)
+
+    # Ensure the specified model exists in the database
+    model_name = args.model.split('/')[-1] if '/' in args.model else args.model
+    model_type = 'huggingface' if '/' in args.model else 'local'
+    ensure_model_exists(model_name, model_type, args.model, 16)  # Default 16GB VRAM
     set_allowed_dir(args.dir)
     set_optimize(args.optimize)
     set_ffmpeg(args.ffmpeg)

vidai/database.py

@@ -559,6 +559,15 @@ def init_db(conn) -> None:
             )
         ''')
 
+    # Insert default model if not exist
+    default_model_name = 'Qwen/Qwen2.5-VL-7B-Instruct'
+    if config['type'] == 'mysql':
+        cursor.execute('INSERT IGNORE INTO models (name, type, path, vram_estimate) VALUES (?, ?, ?, ?)',
+                       ('Qwen 2.5 VL 7B Instruct', 'huggingface', default_model_name, 16))
+    else:
+        cursor.execute('INSERT OR IGNORE INTO models (name, type, path, vram_estimate) VALUES (?, ?, ?, ?)',
+                       ('Qwen 2.5 VL 7B Instruct', 'huggingface', default_model_name, 16))
+
     # Add data column to sessions table if it doesn't exist
     try:
         if config['type'] == 'mysql':
@@ -1941,3 +1950,21 @@ def get_model_by_id(model_id: int) -> Optional[Dict[str, Any]]:
     row = cursor.fetchone()
     conn.close()
     return dict(row) if row else None
+
+
+def ensure_model_exists(name: str, model_type: str, path: str, vram_estimate: int = 0) -> None:
+    """Ensure a model exists in the database, create if not."""
+    conn = get_db_connection()
+    cursor = conn.cursor()
+
+    # Check if model already exists by path
+    cursor.execute('SELECT id FROM models WHERE path = ?', (path,))
+    existing = cursor.fetchone()
+
+    if not existing:
+        # Create the model
+        cursor.execute('INSERT INTO models (name, type, path, vram_estimate) VALUES (?, ?, ?, ?)',
+                       (name, model_type, path, vram_estimate))
+        conn.commit()
+
+    conn.close()
\ No newline at end of file