Mightys/Notebook_Scripts

1# dmagic_swarm.py2"""3Mágica %download para IPython / Jupyter / Colab.4Lee el token de CivitAI desde ~/.civitai_token.pkl (pickle)5Incluye barra de progreso interactiva, renombrado con -o y autogeneración de metadata .swarm.json6 7Uso:8    import download_magic9    %download https://civitai.com/api/download/models/..., https://huggingface.co/...10    %download https://drive.google.com/file/d/123.../view -o MiModelo.safetensors11"""12 13from IPython import get_ipython14from IPython.core.magic import register_line_magic15from IPython.display import HTML, display, clear_output16import ipywidgets as widgets17import os18import subprocess19import re20import requests21import pickle22import json23import base6424from pathlib import Path25 26# ---------- leer token desde pickle ----------27TOKEN_FILE = Path.home() / ".civitai_token.pkl"28token = None29if TOKEN_FILE.exists():30    try:31        token = pickle.loads(TOKEN_FILE.read_bytes())32    except Exception:33        pass34# ---------------------------------------------35 36def ejecutar_con_progreso(cmd, is_gdown=False):37    """Ejecuta el comando leyendo la salida para animar la barra de progreso."""38    progress_bar = widgets.IntProgress(39        value=0, min=0, max=100, 40        description='Progreso:', 41        bar_style='info', 42        orientation='horizontal', 43        layout=widgets.Layout(width='80%')44    )45    status_label = widgets.Label(value="Iniciando descarga...")46    display(widgets.VBox([progress_bar, status_label]))47    48    process = subprocess.Popen(49        cmd, 50        stdout=subprocess.PIPE, 51        stderr=subprocess.STDOUT, 52        text=True, 53        bufsize=1, 54        universal_newlines=True55    )56    57    for line in process.stdout:58        if not is_gdown:59            match_pct = re.search(r'\((\d+)%\)', line)60            if match_pct:61                progress_bar.value = int(match_pct.group(1))62            63            match_speed = re.search(r'DL:([^\s]+)', line)64            if match_speed:65                status_label.value = f"Descargando... Velocidad: {match_speed.group(1)}"66        else:67            match_pct = re.search(r'(\d{1,3})%', line)68            if match_pct:69                progress_bar.value = int(match_pct.group(1))70                status_label.value = "Descargando con gdown..."71 72    process.wait()73    progress_bar.value = 10074    progress_bar.bar_style = 'success'75    status_label.value = "¡Descarga de este archivo completada! ✅"76 77def generar_metadata_civitai(url, dest_path, pretty_name):78    """Descarga información del modelo desde la API de Civitai y crea el .swarm.json"""79    match = re.search(r'models/(\d+)', url)80    if not match:81        match = re.search(r'modelVersionId=(\d+)', url)82    if not match:83        return84 85    version_id = match.group(1)86 87    try:88        display(HTML(f"<p style='color:#4682B4;'>📄 Obteniendo metadata y preview desde Civitai para el ID <code>{version_id}</code>...</p>"))89        90        # Obtener datos de la versión del modelo91        v_resp = requests.get(f"https://civitai.com/api/v1/model-versions/{version_id}", timeout=10)92        if v_resp.status_code != 200:93            display(HTML(f"<p style='color:orange;'>⚠️ No se pudo obtener la metadata (Puede que el modelo esté oculto). Saltando .json...</p>"))94            return95            96        v_data = v_resp.json()97 98        # Obtener datos generales del modelo (para autor y tags)99        m_id = v_data.get("modelId")100        m_data = {}101        if m_id:102            m_resp = requests.get(f"https://civitai.com/api/v1/models/{m_id}", timeout=10)103            if m_resp.status_code == 200:104                m_data = m_resp.json()105 106        # Procesar imagen a base64107        thumbnail_b64 = ""108        images = v_data.get("images", [])109        if images:110            img_url = images[0].get("url")111            if img_url:112                try:113                    i_resp = requests.get(img_url, timeout=10)114                    if i_resp.status_code == 200:115                        content_type = i_resp.headers.get('content-type', 'image/jpeg')116                        b64_str = base64.b64encode(i_resp.content).decode('utf-8')117                        thumbnail_b64 = f"data:{content_type};base64,{b64_str}"118                except Exception:119                    pass120 121        # Construir diccionario de Swarm122        tags = ", ".join(m_data.get("tags", []))123        trigger_words = ", ".join(v_data.get("trainedWords", []))124        model_name = m_data.get("name", v_data.get("name", "Modelo"))125        version_name = v_data.get("name", "")126        127        desc_html = f'<p>From <a href="https://civitai.com/models/{m_id}?modelVersionId={version_id}" target="_blank">Civitai</a></p><hr />'128        desc_html += v_data.get("description", "") or m_data.get("description", "")129 130        swarm_metadata = {131            "modelspec.title": f"{model_name} - {version_name}",132            "modelspec.description": desc_html,133            "modelspec.date": v_data.get("createdAt", ""),134            "modelspec.author": m_data.get("creator", {}).get("username", ""),135            "modelspec.trigger_phrase": trigger_words,136            "modelspec.tags": tags,137            "modelspec.thumbnail": thumbnail_b64,138            "modelspec.usage_hint": v_data.get("baseModel", "")139        }140 141        # Guardar el JSON142        base_name = os.path.splitext(pretty_name)[0]143        if base_name == "Desconocido":144            return145            146        json_filename = f"{base_name}.swarm.json"147        json_path = os.path.join(dest_path, json_filename)148        149        with open(json_path, "w", encoding="utf-8") as f:150            json.dump(swarm_metadata, f, ensure_ascii=False, indent=2)151            152        display(HTML(f"<p style='color:lightgreen;'>✅ Archivo de Swarm <code>{json_filename}</code> creado con éxito.</p>"))153 154    except Exception as e:155        display(HTML(f"<p style='color:orange;'>⚠️ Error construyendo metadata: {e}</p>"))156 157@register_line_magic158def download(line):159    """160    %download url1 [-o nombre1.ext], url2, ...161    Muestra solo el nombre real y respeta la carpeta actual (%cd).162    Auto-detecta enlaces de Drive para usar gdown y permite renombrado.163    """164    dest = os.getcwd()165    # Separar por comas para soportar múltiples descargas en la misma línea166    items = [item.strip() for item in line.split(",") if item.strip()]167 168    for item in items:169        # Extraer URL y el posible nombre personalizado usando regex para atrapar " -o "170        parts = re.split(r'\s+-o\s+', item, maxsplit=1)171        url = parts[0].strip()172        custom_name = parts[1].strip() if len(parts) > 1 else None173 174        # ---------- Google Drive (gdown) ----------175        if "drive.google.com" in url:176            pretty = custom_name if custom_name else "Archivo de Google Drive (Gdown gestiona el nombre)"177            display(HTML(f"<hr><h3 style='color:yellow;'>🛸 Descargando (gdown): <code>{pretty}</code></h3>"178                         f"<h4 style='color:cyan;'>📁 Destino: <code>{dest}</code></h4>"))179            180            cmd = ["gdown", "--fuzzy", url]181            # Si hay nombre custom, le damos la ruta exacta, si no, el directorio182            if custom_name:183                cmd.extend(["-O", os.path.join(dest, custom_name)])184            else:185                cmd.extend(["-O", f"{dest}/"]) 186            187            ejecutar_con_progreso(cmd, is_gdown=True)188 189        # ---------- CivitAI y CivitaiArchive (aria2) ----------190        elif "civitai.com" in url or "civitaiarchive.com" in url:191            if not token and "civitai.com" in url:192                display(HTML("<h4 style='color:red;'>⚠️ Token de CivitAI no encontrado (algunos modelos pueden fallar).</h4>"))193            194            url_token = url195            if token and ("civitai.com" in url or "civitaiarchive.com" in url):196                url_token = f"{url}{'&' if '?' in url else '?'}token={token}"197 198            # Si el usuario mandó nombre manual, lo usamos. Si no, consultamos la API para saberlo.199            if custom_name:200                pretty = custom_name201            else:202                try:203                    with requests.get(url_token, stream=True, timeout=5) as r:204                        cd = r.headers.get('Content-Disposition', '')205                        match = re.findall(r'filename[*]?=(?:UTF-8\'\')?["\']?([^"\';]+)["\']?', cd)206                        pretty = match[0] if match else url.split('/')[-1].split('?')[0]207                except Exception:208                    pretty = url.split('/')[-1].split('?')[0]209 210            display(HTML(f"<hr><h3 style='color:yellow;'>📥 Descargando (aria2): <code>{pretty}</code></h3>"211                         f"<h4 style='color:cyan;'>📁 Destino: <code>{dest}</code></h4>"))212 213            cmd = [214                "aria2c", "--summary-interval=1",215                "-c", "-x", "16", "-s", "16", "-k", "1M",216                "-d", dest217            ]218            219            # Si hay nombre personalizado, usamos -o. Si no, usamos content-disposition220            if custom_name:221                cmd.extend(["-o", custom_name])222            else:223                cmd.append("--content-disposition")224                225            cmd.append(url_token)226            ejecutar_con_progreso(cmd, is_gdown=False)227 228            # Generar metadata de Civitai al finalizar229            if pretty and pretty != "Desconocido":230                generar_metadata_civitai(url, dest, pretty)231 232        # ---------- HuggingFace / Otros (aria2) ----------233        else:234            pretty = custom_name if custom_name else url.split('/')[-1].split('?')[0]235            display(HTML(f"<hr><h3 style='color:yellow;'>📥 Descargando (aria2): <code>{pretty}</code></h3>"236                         f"<h4 style='color:cyan;'>📁 Destino: <code>{dest}</code></h4>"))237            238            cmd = [239                "aria2c", "--summary-interval=1",240                "-c", "-x", "16", "-s", "16", "-k", "1M",241                "-d", dest242            ]243            244            if custom_name:245                cmd.extend(["-o", custom_name])246            else:247                cmd.extend(["-o", pretty])248                249            cmd.append(url)250            ejecutar_con_progreso(cmd, is_gdown=False)251 252            # Limpiar hashes de 64 caracteres típicos de HuggingFace si quedó suelto y no hay nombre custom253            if not custom_name and ("huggingface.co" in url):254                for f in os.listdir(dest):255                    if re.fullmatch(r'[0-9a-f]{64}', f):256                        os.rename(os.path.join(dest, f), os.path.join(dest, pretty))257                        break258 259    display(HTML("<br><h3 style='color:lightgreen;'>✅ ¡Todas las descargas solicitadas han finalizado!</h3>"))260 261# Registramos la línea mágica al importar el módulo262get_ipython().register_magic_function(download, magic_kind='line')