Delete model_downloader.py
This commit is contained in:
Iván Eduardo Chavez Ayub
committed by
GitHub
GitHub
parent
555d1e24c3
commit
a75a73f041
@@ -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!")
|
||||
Reference in New Issue
Block a user