import os import sys import asyncio import aiohttp import urllib.parse import argparse import shutil import random import hashlib from tqdm.asyncio import tqdm import nest_asyncio # Apply nest_asyncio globally nest_asyncio.apply() # --- STEALTH AGENTS --- USER_AGENTS = [ "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2 Safari/605.1.15", "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:123.0) Gecko/20100101 Firefox/123.0", "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" ] class Speedster: def __init__(self, num_threads=16, chunk_size_mb=2, max_retries=5): self.num_threads = num_threads self.chunk_size = chunk_size_mb * 1024 * 1024 self.max_retries = max_retries self.is_windows = os.name == 'nt' def _get_headers(self, token=None): # Stealth Mode: Pick random agent headers = { "User-Agent": random.choice(USER_AGENTS), "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8", "Accept-Language": "en-US,en;q=0.5", } # Priority: Function Arg > Env Var final_token = token or os.environ.get("CIVITAI_TOKEN") or os.environ.get("HF_TOKEN") if final_token: headers["Authorization"] = f"Bearer {final_token}" return headers def _write_data(self, fd, data, offset): if self.is_windows: os.lseek(fd, offset, 0) os.write(fd, data) else: os.pwrite(fd, data, offset) def _check_disk_space(self, path, required_bytes): total, used, free = shutil.disk_usage(os.path.dirname(os.path.abspath(path)) or '.') if free < required_bytes: raise OSError(f"❌ Disk Full! Required: {required_bytes/1024**3:.2f}GB, Free: {free/1024**3:.2f}GB") async def _fetch_chunk(self, session, url, start, end, fd, headers, pbar): chunk_headers = headers.copy() chunk_headers['Range'] = f'bytes={start}-{end}' for attempt in range(self.max_retries): try: async with session.get(url, headers=chunk_headers, timeout=60) as response: response.raise_for_status() offset = start async for chunk in response.content.iter_chunked(self.chunk_size): if chunk: self._write_data(fd, chunk, offset) offset += len(chunk) pbar.update(len(chunk)) return except Exception as e: if attempt < self.max_retries - 1: await asyncio.sleep(1 * (attempt + 1)) else: # tqdm.write(f"❌ Chunk failed: {e}") # Silent fail to reduce spam raise e async def _download_logic(self, url, dest_path, token=None, extract=False, verify=False): headers = self._get_headers(token) async with aiohttp.ClientSession() as session: # 1. Metadata & Smart Filename try: async with session.get(url, headers=headers, allow_redirects=True) as response: final_url = response.url total_size = int(response.headers.get('Content-Length', 0)) filename = "downloaded_file.bin" if os.path.isdir(dest_path): if 'Content-Disposition' in response.headers: try: cd = response.headers['Content-Disposition'] filename = cd.split('filename=')[1].strip('"') except: filename = os.path.basename(urllib.parse.urlparse(str(final_url)).path) else: path_filename = os.path.basename(urllib.parse.urlparse(str(final_url)).path) if path_filename: filename = path_filename filename = filename.split('?')[0] final_path = os.path.join(dest_path, filename) else: final_path = dest_path # 2. Smart Skip (File Exists & Size Match) if os.path.exists(final_path): local_size = os.path.getsize(final_path) if local_size == total_size and total_size > 0: print(f"⏩ [Skip] {os.path.basename(final_path)} exists & valid.") return final_path # 3. Disk Check self._check_disk_space(dest_path, total_size) print(f"🚀 [Speedster] {os.path.basename(final_path)}") print(f"📦 Size: {total_size / (1024**3):.2f} GB | Threads: {self.num_threads}") # 4. Atomic Download (.part file) part_path = final_path + ".part" fd = os.open(part_path, os.O_RDWR | os.O_CREAT | os.O_BINARY if self.is_windows else os.O_RDWR | os.O_CREAT) os.ftruncate(fd, total_size) chunk_size = total_size // self.num_threads chunks = [] for i in range(self.num_threads): start = i * chunk_size end = total_size - 1 if i == self.num_threads - 1 else (start + chunk_size - 1) chunks.append((start, end)) pbar = tqdm(total=total_size, unit='iB', unit_scale=True, unit_divisor=1024, desc="⚡ DL") tasks = [asyncio.create_task(self._fetch_chunk(session, final_url, s, e, fd, headers, pbar)) for s, e in chunks] await asyncio.gather(*tasks) pbar.close() os.close(fd) # 5. Finalize (Rename .part -> .safetensors) if os.path.exists(final_path): os.remove(final_path) os.rename(part_path, final_path) print("✅ Download Complete!") # 6. Auto-Extract if extract: print(f"📂 Extracting {filename}...") try: shutil.unpack_archive(final_path, dest_path) print("✅ Extracted successfully.") except Exception as e: print(f"⚠️ Extraction failed: {e}") return final_path except Exception as e: print(f"❌ Error: {e}") if 'part_path' in locals() and os.path.exists(part_path): os.remove(part_path) # Cleanup corrupt partial return None def download(self, url, dest_path=".", token=None, extract=False, verify=False): loop = asyncio.get_event_loop() if loop.is_running(): nest_asyncio.apply() return loop.run_until_complete(self._download_logic(url, dest_path, token, extract, verify)) else: return asyncio.run(self._download_logic(url, dest_path, token, extract, verify)) # --- CLI HANDLER --- def cli_main(): parser = argparse.ArgumentParser(description="Speedster v16.0.3 - Juggernaut Update") parser.add_argument("input", help="URL file OR path to .txt file for batch download") parser.add_argument("--out", "-o", default=".", help="Output folder") parser.add_argument("--token", "-t", help="Auth Token (Optional if env var set)") parser.add_argument("--threads", type=int, default=16, help="Thread count") parser.add_argument("--extract", "-x", action="store_true", help="Auto extract archives") parser.add_argument("--batch", "-b", action="store_true", help="Force treat input as batch file") args = parser.parse_args() engine = Speedster(num_threads=args.threads) # BATCH MODE LOGIC urls = [] if args.batch or (os.path.isfile(args.input) and args.input.endswith(".txt")): print(f"📜 Batch Mode: Reading {args.input}") with open(args.input, 'r') as f: urls = [line.strip() for line in f if line.strip() and not line.startswith("#")] else: urls = [args.input] for i, url in enumerate(urls): if len(urls) > 1: print(f"\n--- Processing {i+1}/{len(urls)} ---") engine.download(url, args.out, args.token, extract=args.extract) if __name__ == "__main__": cli_main()