""" AI Model Downloader for Warlock-Studio Descarga automática de modelos AI cuando no están presentes """ import os import sys import requests import zipfile from pathlib import Path from typing import Optional import tempfile import shutil from tkinter import messagebox import threading import time class ModelDownloader: def __init__(self): self.base_path = self._get_base_path() self.ai_models_path = os.path.join(self.base_path, "AI-onnx") self.download_url = "https://github.com/Ivan-Ayub97/Warlock-Studio/releases/download/v4.0/AI-onnx-models.zip" self.backup_urls = [ "https://sourceforge.net/projects/warlock-studio/files/AI-Models/AI-onnx-models.zip", # Agregar más URLs de respaldo aquí ] def _get_base_path(self) -> str: """Obtener la ruta base de la aplicación""" if getattr(sys, '_MEIPASS', None): return sys._MEIPASS return os.path.dirname(os.path.abspath(__file__)) def check_models_exist(self) -> bool: """Verificar si los modelos AI existen""" if not os.path.exists(self.ai_models_path): return False required_models = [ "BSRGANx2_fp16.onnx", "BSRGANx4_fp16.onnx", "GFPGANv1.4.fp16.onnx", "IRCNN_Lx1_fp16.onnx", "IRCNN_Mx1_fp16.onnx", "RIFE_Lite_fp32.onnx", "RIFE_fp32.onnx", "RealESRGANx4_fp16.onnx", "RealESRNetx4_fp16.onnx", "RealESR_Animex4_fp16.onnx", "RealESR_Gx4_fp16.onnx", "RealSRx4_Anime_fp16.onnx", "super-resolution-10.onnx" ] for model in required_models: model_path = os.path.join(self.ai_models_path, model) if not os.path.exists(model_path): return False return True def download_with_progress(self, url: str, destination: str, progress_callback=None) -> bool: """Descargar archivo con barra de progreso""" try: response = requests.get(url, stream=True) response.raise_for_status() total_size = int(response.headers.get('content-length', 0)) downloaded = 0 with open(destination, 'wb') as file: for chunk in response.iter_content(chunk_size=8192): if chunk: file.write(chunk) downloaded += len(chunk) if progress_callback and total_size > 0: progress = (downloaded / total_size) * 100 progress_callback(progress, downloaded, total_size) return True except Exception as e: print(f"Error downloading from {url}: {str(e)}") return False def extract_zip(self, zip_path: str, extract_path: str) -> bool: """Extraer archivo ZIP""" try: with zipfile.ZipFile(zip_path, 'r') as zip_ref: zip_ref.extractall(extract_path) return True except Exception as e: print(f"Error extracting {zip_path}: {str(e)}") return False def download_models(self, progress_callback=None) -> bool: """Descargar modelos AI""" if self.check_models_exist(): print("AI models already exist, skipping download") return True print("AI models not found, downloading...") # Crear directorio temporal with tempfile.TemporaryDirectory() as temp_dir: zip_path = os.path.join(temp_dir, "AI-onnx-models.zip") # Intentar descargar desde URL principal success = False for url in [self.download_url] + self.backup_urls: print(f"Attempting download from: {url}") if self.download_with_progress(url, zip_path, progress_callback): success = True break else: print(f"Failed to download from {url}, trying next URL...") if not success: return False # Crear directorio de destino os.makedirs(self.ai_models_path, exist_ok=True) # Extraer archivos if self.extract_zip(zip_path, self.base_path): print("AI models downloaded and extracted successfully") return True else: return False def download_models_async(self, completion_callback=None, progress_callback=None): """Descargar modelos de forma asíncrona""" def download_thread(): success = self.download_models(progress_callback) if completion_callback: completion_callback(success) thread = threading.Thread(target=download_thread) thread.daemon = True thread.start() return thread # Función para integrar con el código principal def ensure_models_available(show_dialog=True) -> bool: """Asegurar que los modelos AI estén disponibles""" downloader = ModelDownloader() if downloader.check_models_exist(): return True if show_dialog: from tkinter import messagebox result = messagebox.askyesno( "AI Models Required", "AI models are required but not found. Would you like to download them now?\n\n" "This will download approximately 327MB of data.\n\n" "Click 'Yes' to download or 'No' to continue without AI functionality." ) if not result: return False # Mostrar ventana de progreso simple try: import tkinter as tk from tkinter import ttk progress_window = tk.Toplevel() progress_window.title("Downloading AI Models") progress_window.geometry("400x120") progress_window.resizable(False, False) progress_label = tk.Label(progress_window, text="Downloading AI models...") progress_label.pack(pady=10) progress_bar = ttk.Progressbar(progress_window, length=300, mode='determinate') progress_bar.pack(pady=10) status_label = tk.Label(progress_window, text="Starting download...") status_label.pack(pady=5) download_complete = [False] def update_progress(percentage, downloaded, total): progress_bar['value'] = percentage mb_downloaded = downloaded / (1024 * 1024) mb_total = total / (1024 * 1024) status_label.config(text=f"Downloaded: {mb_downloaded:.1f} MB / {mb_total:.1f} MB") progress_window.update() def on_completion(success): download_complete[0] = True if success: status_label.config(text="Download completed successfully!") else: status_label.config(text="Download failed!") progress_window.after(2000, progress_window.destroy) # Iniciar descarga downloader.download_models_async(on_completion, update_progress) # Mantener ventana activa hasta completar while not download_complete[0]: progress_window.update() time.sleep(0.1) return download_complete[0] except ImportError: # Si no hay tkinter, descargar sin interfaz gráfica return downloader.download_models() if __name__ == "__main__": downloader = ModelDownloader() if not downloader.check_models_exist(): print("Downloading AI models...") success = downloader.download_models() if success: print("Models downloaded successfully!") else: print("Failed to download models!") else: print("AI models already available!")