import os import uuid import requests from pathlib import Path from typing import Optional, Union from fastapi import UploadFile # ============================== # STORAGE CONFIG # ============================== BASE_DIR = Path(__file__).resolve().parents[1] UPLOAD_DIR = str(BASE_DIR / "jobs" / "uploads") os.makedirs(UPLOAD_DIR, exist_ok=True) CHUNK_SIZE = 1024 * 1024 # 1MB streaming for large media # ============================== # CORE RESOLVER # ============================== def resolve_input( source: Optional[str] = None, upload: Optional[UploadFile] = None, raw_bytes: Optional[bytes] = None ) -> str: """ Universal ingestion layer for all pipeline systems. Supports: - UploadFile (FastAPI / Gradio) - URL download (http/https) - Local filesystem path - Raw bytes input (future automation nodes) """ # ---------------------------------------- # CASE 1: UploadFile (Gradio / FastAPI) # ---------------------------------------- if upload is not None: filename = f"{uuid.uuid4()}_{upload.filename or 'upload.mp4'}" path = os.path.join(UPLOAD_DIR, filename) with open(path, "wb") as f: while True: chunk = upload.file.read(CHUNK_SIZE) if not chunk: break f.write(chunk) return path # ---------------------------------------- # CASE 2: Raw bytes (automation / webhook) # ---------------------------------------- if raw_bytes is not None: filename = f"{uuid.uuid4()}.mp4" path = os.path.join(UPLOAD_DIR, filename) with open(path, "wb") as f: f.write(raw_bytes) return path # ---------------------------------------- # CASE 3: URL input (YouTube, TikTok, direct mp4) # ---------------------------------------- if source and source.startswith(("http://", "https://")): filename = f"{uuid.uuid4()}.mp4" path = os.path.join(UPLOAD_DIR, filename) headers = { "User-Agent": "Mozilla/5.0 (compatible; BasyxBot/1.0)" } with requests.get(source, stream=True, headers=headers, timeout=60) as r: r.raise_for_status() with open(path, "wb") as f: for chunk in r.iter_content(chunk_size=CHUNK_SIZE): if chunk: f.write(chunk) return path # ---------------------------------------- # CASE 4: Local file path # ---------------------------------------- if source and os.path.exists(source): return source # ---------------------------------------- # INVALID INPUT HANDLING # ---------------------------------------- raise ValueError( "resolve_input failed: no valid source, upload, or raw_bytes provided" ) # ============================== # OPTIONAL HELPERS (V11 READY) # ============================== def detect_input_type(source: str) -> str: """ Lightweight classifier for routing decisions upstream. """ if source.startswith(("http://", "https://")): return "url" if os.path.exists(source): return "file" return "unknown" def normalize_source(source: str) -> str: """ Cleans input strings for downstream consistency. """ if not source: return source return source.strip()