diff --git a/api/__init__.py b/api/__init__.py index 4ee2bb73..a3ae903b 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -27,6 +27,7 @@ from tqdm import tqdm import traceback import sys from lollms.console import MainMenu +import urllib __author__ = "parisneo" __github__ = "https://github.com/ParisNeo/lollms-webui" @@ -142,6 +143,7 @@ class LoLLMsAPPI(): # This is used to keep track of messages self.full_message_list = [] self.current_room_id = None + self.download_infos={} # ========================================================================================= # Socket IO stuff # ========================================================================================= @@ -158,10 +160,31 @@ class LoLLMsAPPI(): self.cancel_gen = True ASCIIColors.error(f'Client {request.sid} canceled generation') - + @socketio.on('cancel_install') + def cancel_install(data): + model_name = data["model_name"] + binding_folder = data["binding_folder"] + model_url = data["model_url"] + signature = f"{model_name}_{binding_folder}_{model_url}" + self.download_infos[signature]["cancel"]=True + + @socketio.on('install_model') def install_model(data): room_id = request.sid + + def get_file_size(url): + # Send a HEAD request to retrieve file metadata + response = urllib.request.urlopen(url) + + # Extract the Content-Length header value + file_size = response.headers.get('Content-Length') + + # Convert the file size to integer + if file_size: + file_size = int(file_size) + + return file_size def install_model_(): print("Install model triggered") model_path = data["path"] @@ -175,6 +198,16 @@ class LoLLMsAPPI(): model_name = filename binding_folder = self.config["binding_name"] model_url = model_path + signature = f"{model_name}_{binding_folder}_{model_url}" + self.download_infos[signature]={ + "start_time":datetime.now(), + "total_size":get_file_size(model_path), + "downloaded_size":0, + "progress":0, + "speed":0, + "cancel":False + } + if installation_path.exists(): print("Error: Model already exists") socketio.emit('install_progress',{ @@ -182,38 +215,111 @@ class LoLLMsAPPI(): 'error': 'model already exists', 'model_name' : model_name, 'binding_folder' : binding_folder, - 'model_url' : model_url + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': self.download_infos[signature]['progress'], + 'speed': self.download_infos[signature]['speed'], }, room=room_id ) socketio.emit('install_progress',{ - 'status': 'progress', + 'status': True, 'progress': progress, 'model_name' : model_name, 'binding_folder' : binding_folder, - 'model_url' : model_url + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': self.download_infos[signature]['progress'], + 'speed': self.download_infos[signature]['speed'], + }, room=room_id) - def callback(progress): - socketio.emit('install_progress',{ - 'status': 'progress', - 'progress': progress, - 'model_name' : model_name, - 'binding_folder' : binding_folder, - 'model_url' : model_url - }, room=room_id) + def callback(downloaded_size, total_size): + progress = (downloaded_size / total_size) * 100 + now = datetime.now() + dt = (now - self.download_infos[signature]['start_time']).total_seconds() + speed = downloaded_size/dt + self.download_infos[signature]['downloaded_size'] = downloaded_size + self.download_infos[signature]['speed'] = speed + + if self.download_infos[signature]['progress']-progress>2: + self.download_infos[signature]['progress'] = progress + socketio.emit('install_progress',{ + 'status': True, + 'model_name' : model_name, + 'binding_folder' : binding_folder, + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': self.download_infos[signature]['progress'], + 'speed': self.download_infos[signature]['speed'], + }, room=room_id) + + if self.download_infos[signature]["cancel"]: + raise Exception("canceled") if hasattr(self.binding, "download_model"): - self.binding.download_model(model_path, installation_path, callback) + try: + self.binding.download_model(model_path, installation_path, callback) + except Exception as ex: + ASCIIColors.warning(str(ex)) + socketio.emit('install_progress',{ + 'status': False, + 'error': 'canceled', + 'model_name' : model_name, + 'binding_folder' : binding_folder, + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': self.download_infos[signature]['progress'], + 'speed': self.download_infos[signature]['speed'], + }, room=room_id + ) + del self.download_infos[signature] + installation_path.unlink() + return + else: - self.download_file(model_path, installation_path, callback) + try: + self.download_file(model_path, installation_path, callback) + except Exception as ex: + ASCIIColors.warning(str(ex)) + socketio.emit('install_progress',{ + 'status': False, + 'error': 'canceled', + 'model_name' : model_name, + 'binding_folder' : binding_folder, + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': self.download_infos[signature]['progress'], + 'speed': self.download_infos[signature]['speed'], + }, room=room_id + ) + del self.download_infos[signature] + installation_path.unlink() + return socketio.emit('install_progress',{ 'status': True, 'error': '', 'model_name' : model_name, 'binding_folder' : binding_folder, - 'model_url' : model_url + 'model_url' : model_url, + 'start_time': self.download_infos[signature]['start_time'], + 'total_size': self.download_infos[signature]['total_size'], + 'downloaded_size': self.download_infos[signature]['downloaded_size'], + 'progress': 100, + 'speed': self.download_infos[signature]['speed'], }, room=room_id) + del self.download_infos[signature] + tpe = threading.Thread(target=install_model_, args=()) tpe.start() @@ -223,12 +329,25 @@ class LoLLMsAPPI(): installation_dir = self.lollms_paths.personal_models_path/self.config["binding_name"] filename = Path(model_path).name installation_path = installation_dir / filename + + model_name = filename + binding_folder = self.config["binding_name"] if not installation_path.exists(): - socketio.emit('install_progress',{'status': False, 'error': 'The model does not exist', 'url':model_path}, room=request.sid) + socketio.emit('install_progress',{ + 'status': False, + 'error': 'The model does not exist', + 'model_name' : model_name, + 'binding_folder' : binding_folder + }, room=request.sid) installation_path.unlink() - socketio.emit('install_progress',{'status': True, 'error': '', 'url':model_path}, room=request.sid) + socketio.emit('install_progress',{ + 'status': True, + 'error': '', + 'model_name' : model_name, + 'binding_folder' : binding_folder + }, room=request.sid) @socketio.on('upload_file') @@ -469,8 +588,7 @@ class LoLLMsAPPI(): file.write(chunk) downloaded_size += len(chunk) if callback is not None: - percentage = (downloaded_size / total_size) * 100 - callback(percentage) + callback(downloaded_size, total_size) progress_bar.update(len(chunk)) if callback is not None: