From 5387576c59632f6bc85c81de4d19219f63437d0a Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 15 Mar 2023 15:11:04 -0400 Subject: api error handler --- modules/api/api.py | 45 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) (limited to 'modules/api/api.py') diff --git a/modules/api/api.py b/modules/api/api.py index 35e17afc..8d4ff4e0 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -6,8 +6,11 @@ import uvicorn from threading import Lock from io import BytesIO from gradio.processing_utils import decode_base64_to_file -from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request, Response +from fastapi import APIRouter, Depends, FastAPI, Request, Response from fastapi.security import HTTPBasic, HTTPBasicCredentials +from fastapi.exceptions import HTTPException +from fastapi.responses import JSONResponse +from fastapi.encoders import jsonable_encoder from secrets import compare_digest import modules.shared as shared @@ -90,6 +93,16 @@ def encode_pil_to_base64(image): return base64.b64encode(bytes_data) def api_middleware(app: FastAPI): + rich_available = True + try: + import anyio # importing just so it can be placed on silent list + import starlette # importing just so it can be placed on silent list + from rich.console import Console + console = Console() + except: + import traceback + rich_available = False + @app.middleware("http") async def log_and_time(req: Request, call_next): ts = time.time() @@ -110,6 +123,36 @@ def api_middleware(app: FastAPI): )) return res + def handle_exception(request: Request, e: Exception): + err = { + "error": type(e).__name__, + "detail": vars(e).get('detail', ''), + "body": vars(e).get('body', ''), + "errors": str(e), + } + print(f"API error: {request.method}: {request.url} {err}") + if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions + if rich_available: + console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200])) + else: + traceback.print_exc() + return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err)) + + @app.middleware("http") + async def exception_handling(request: Request, call_next): + try: + return await call_next(request) + except Exception as e: + return handle_exception(request, e) + + @app.exception_handler(Exception) + async def fastapi_exception_handler(request: Request, e: Exception): + return handle_exception(request, e) + + @app.exception_handler(HTTPException) + async def http_exception_handler(request: Request, e: HTTPException): + return handle_exception(request, e) + class Api: def __init__(self, app: FastAPI, queue_lock: Lock): -- cgit v1.2.3 From 4cbbb881ee530d9b9ba18027e2b0057e6a2c4ee1 Mon Sep 17 00:00:00 2001 From: Φφ <42910943+Brawlence@users.noreply.github.com> Date: Thu, 9 Mar 2023 07:56:19 +0300 Subject: Unload checkpoints on Request MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit …to free VRAM. New Action buttons in the settings to manually free and reload checkpoints, essentially juggling models between RAM and VRAM. --- modules/api/api.py | 14 +++++++++++++- modules/sd_models.py | 22 +++++++++++++++++++++- modules/ui.py | 22 ++++++++++++++++++++++ 3 files changed, 56 insertions(+), 2 deletions(-) (limited to 'modules/api/api.py') diff --git a/modules/api/api.py b/modules/api/api.py index 35e17afc..f52f7fef 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -18,7 +18,7 @@ from modules.textual_inversion.textual_inversion import create_embedding, train_ from modules.textual_inversion.preprocess import preprocess from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork from PIL import PngImagePlugin,Image -from modules.sd_models import checkpoints_list +from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights from modules.sd_models_config import find_checkpoint_config_near_filename from modules.realesrgan_model import get_realesrgan_models from modules import devices @@ -150,6 +150,8 @@ class Api: self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=TrainResponse) self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=TrainResponse) self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=MemoryResponse) + self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"]) + self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"]) self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=ScriptsList) def add_api_route(self, path: str, endpoint, **kwargs): @@ -412,6 +414,16 @@ class Api: return {} + def unloadapi(self): + unload_model_weights() + + return {} + + def reloadapi(self): + reload_model_weights() + + return {} + def skip(self): shared.state.skip() diff --git a/modules/sd_models.py b/modules/sd_models.py index f0cb1240..f9dd0521 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -494,7 +494,7 @@ def reload_model_weights(sd_model=None, info=None): if sd_model is None or checkpoint_config != sd_model.used_config: del sd_model checkpoints_loaded.clear() - load_model(checkpoint_info, already_loaded_state_dict=state_dict, time_taken_to_load_state_dict=timer.records["load weights from disk"]) + load_model(checkpoint_info, already_loaded_state_dict=state_dict) return shared.sd_model try: @@ -517,3 +517,23 @@ def reload_model_weights(sd_model=None, info=None): print(f"Weights loaded in {timer.summary()}.") return sd_model + +def unload_model_weights(sd_model=None, info=None): + from modules import lowvram, devices, sd_hijack + timer = Timer() + + if shared.sd_model: + + # shared.sd_model.cond_stage_model.to(devices.cpu) + # shared.sd_model.first_stage_model.to(devices.cpu) + shared.sd_model.to(devices.cpu) + sd_hijack.model_hijack.undo_hijack(shared.sd_model) + shared.sd_model = None + sd_model = None + gc.collect() + devices.torch_gc() + torch.cuda.empty_cache() + + print(f"Unloaded weights {timer.summary()}.") + + return sd_model \ No newline at end of file diff --git a/modules/ui.py b/modules/ui.py index 7e603332..d93ef134 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1491,11 +1491,33 @@ def create_ui(): request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications") download_localization = gr.Button(value='Download localization template', elem_id="download_localization") reload_script_bodies = gr.Button(value='Reload custom script bodies (No ui updates, No restart)', variant='secondary', elem_id="settings_reload_script_bodies") + with gr.Row(): + unload_sd_model = gr.Button(value='Unload SD checkpoint to free VRAM', elem_id="sett_unload_sd_model") + reload_sd_model = gr.Button(value='Reload the last SD checkpoint back into VRAM', elem_id="sett_reload_sd_model") with gr.TabItem("Licenses"): gr.HTML(shared.html("licenses.html"), elem_id="licenses") gr.Button(value="Show all pages", elem_id="settings_show_all_pages") + + + def unload_sd_weights(): + modules.sd_models.unload_model_weights() + + def reload_sd_weights(): + modules.sd_models.reload_model_weights() + + unload_sd_model.click( + fn=unload_sd_weights, + inputs=[], + outputs=[] + ) + + reload_sd_model.click( + fn=reload_sd_weights, + inputs=[], + outputs=[] + ) request_notifications.click( fn=lambda: None, -- cgit v1.2.3