diff --git a/model_downloader.py b/model_downloader.py deleted file mode 100644 index 8954c59..0000000 --- a/model_downloader.py +++ /dev/null @@ -1,224 +0,0 @@ -""" -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!")