github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-02-10 11:34:57 -05:00
Code Issues Packages Projects Releases Wiki Activity

Model Manager Backend Implementation

Browse Source
This commit is contained in:
blessedcoolant
2022-12-25 09:28:46 +13:00
parent e66b1a685c
commit 3521557541
2 changed files with 106 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
backend/invoke_ai_web_server.py
View File

@@ -9,6 +9,7 @@ import io
import base64
import os
import json
import tkinter as tk
from werkzeug.utils import secure_filename
from flask import Flask, redirect, send_from_directory, request, make_response
@@ -17,6 +18,7 @@ from PIL import Image, ImageOps
from PIL.Image import Image as ImageType
from uuid import uuid4
from threading import Event
from tkinter import filedialog
from ldm.generate import Generate
from ldm.invoke.args import Args, APP_ID, APP_VERSION, calculate_init_img_hash
@@ -297,6 +299,87 @@ class InvokeAIWebServer:
config["infill_methods"] = infill_methods()
socketio.emit("systemConfig", config)
@socketio.on('searchForModels')
def handle_search_models():
try:
# Using tkinter to get the filepath because JS doesn't allow
root = tk.Tk()
root.iconify() # for macos
root.withdraw()
root.wm_attributes('-topmost', 1)
root.focus_force()
search_folder = filedialog.askdirectory(parent=root, title='Select Checkpoint Folder')
root.destroy()
if not search_folder:
socketio.emit(
"foundModels",
{'search_folder': None, 'found_models': None},
)
else:
search_folder, found_models = self.generate.model_cache.search_models(search_folder)
socketio.emit(
"foundModels",
{'search_folder': search_folder, 'found_models': found_models},
)
except Exception as e:
self.socketio.emit("error", {"message": (str(e))})
print("\n")
traceback.print_exc()
print("\n")
@socketio.on("addNewModel")
def handle_add_model(new_model_config: dict):
try:
model_name = new_model_config['name']
del new_model_config['name']
model_attributes = new_model_config
update = False
current_model_list = self.generate.model_cache.list_models()
if model_name in current_model_list:
update = True
print(f">> Adding New Model: {model_name}")
self.generate.model_cache.add_model(
model_name=model_name, model_attributes=model_attributes, clobber=True)
self.generate.model_cache.commit(opt.conf)
new_model_list = self.generate.model_cache.list_models()
socketio.emit(
"newModelAdded",
{"new_model_name": model_name,
"model_list": new_model_list, 'update': update},
)
print(f">> New Model Added: {model_name}")
except Exception as e:
self.socketio.emit("error", {"message": (str(e))})
print("\n")
traceback.print_exc()
print("\n")
@socketio.on("deleteModel")
def handle_delete_model(model_name: str):
try:
print(f">> Deleting Model: {model_name}")
self.generate.model_cache.del_model(model_name)
self.generate.model_cache.commit(opt.conf)
updated_model_list = self.generate.model_cache.list_models()
socketio.emit(
"modelDeleted",
{"deleted_model_name": model_name,
"model_list": updated_model_list},
)
print(f">> Model Deleted: {model_name}")
except Exception as e:
self.socketio.emit("error", {"message": (str(e))})
print("\n")
traceback.print_exc()
print("\n")
@socketio.on("requestModelChange")
def handle_set_model(model_name: str):
try: