mirror of
https://github.com/vladmandic/automatic
synced 2026-08-25 22:20:46 +02:00
CalamitousFelicitousness
b12cbcb523
Add a model section to /sdapi/v1/memory with loaded-model bytes summed per pipeline component and device, so clients can tell resident weights from offloaded ones and loop-critical components from edge ones. - walk components over parameters and buffers, dedupe shared storages, key by component name then device type - read the raw model slot so a memory poll never triggers a model load - section is exception-isolated like ram and cuda; reports an error string if the walk races a reload
231 lines
9.5 KiB
Python
231 lines
9.5 KiB
Python
import os
|
|
import time
|
|
from pathlib import Path
|
|
from fastapi import Request, Depends, BackgroundTasks, Response
|
|
from fastapi.exceptions import HTTPException
|
|
from fastapi.responses import FileResponse
|
|
import installer
|
|
from modules import shared
|
|
from modules.logger import log
|
|
from modules.api import models, helpers
|
|
|
|
|
|
def get_js(request: Request):
|
|
file = request.query_params.get("file", None)
|
|
if (file is None) or (len(file) == 0):
|
|
raise HTTPException(status_code=400, detail="file parameter is required")
|
|
# Security: validate path is within allowed directories
|
|
if shared.demo is None:
|
|
raise HTTPException(status_code=503, detail="server not ready")
|
|
allowed_dirs = shared.demo.allowed_paths
|
|
if not any(Path(folder).absolute() in Path(file).absolute().parents for folder in allowed_dirs):
|
|
raise HTTPException(status_code=403, detail=f"file {file}: must be in one of allowed directories")
|
|
ext = file.split('.')[-1]
|
|
if ext not in ['js', 'css', 'map', 'html', 'wasm', 'ttf', 'mjs', 'json']:
|
|
raise HTTPException(status_code=400, detail=f"invalid file extension: {ext}")
|
|
if not os.path.exists(file):
|
|
log.error(f"API: file not found: {file}")
|
|
raise HTTPException(status_code=404, detail=f"file not found: {file}")
|
|
if ext in ['js', 'mjs']:
|
|
media_type = 'application/javascript'
|
|
elif ext in ['map', 'json']:
|
|
media_type = 'application/json'
|
|
elif ext in ['css']:
|
|
media_type = 'text/css'
|
|
elif ext in ['html']:
|
|
media_type = 'text/html'
|
|
elif ext in ['wasm']:
|
|
media_type = 'application/wasm'
|
|
elif ext in ['ttf']:
|
|
media_type = 'font/ttf'
|
|
else:
|
|
media_type = 'application/octet-stream'
|
|
return FileResponse(file, media_type=media_type)
|
|
|
|
def get_version():
|
|
return installer.get_version()
|
|
|
|
def get_icon():
|
|
icon_path = os.path.join(shared.script_path, "ui", "assets", "favicon.png")
|
|
return FileResponse(icon_path, media_type="image/png")
|
|
|
|
def get_manifest():
|
|
from modules import paths
|
|
manifest_path = os.path.join(paths.script_path, "ui", "manifest", "manifest.json")
|
|
log.debug(f"API manifest={manifest_path}")
|
|
return FileResponse(manifest_path, media_type="application/json")
|
|
|
|
def get_motd():
|
|
import requests
|
|
motd = ""
|
|
ver = get_version()
|
|
if ver.get("updated", None) is not None:
|
|
motd = f"version <b>{ver['commit']} {ver['updated']}</b> <span style='color: var(--primary-500)'>{ver['url'].split('/')[-1]}</span><br>" # pylint: disable=use-maxsplit-arg
|
|
if shared.opts.motd:
|
|
try:
|
|
res = requests.get("https://vladmandic.github.io/sdnext/motd", timeout=3)
|
|
if res.status_code == 200:
|
|
msg = (res.text or "").strip()
|
|
log.info(f"MOTD: {msg if len(msg) > 0 else 'N/A'}")
|
|
motd += res.text
|
|
else:
|
|
log.error(f"MOTD: {res.status_code}")
|
|
except Exception as err:
|
|
log.error(f"MOTD: {err}")
|
|
return motd
|
|
|
|
def get_platform():
|
|
from modules.loader import get_packages as loader_get_packages
|
|
return { **installer.get_platform(), **loader_get_packages() }
|
|
|
|
def get_torch():
|
|
return dict(installer.torch_info)
|
|
|
|
def get_log(req: models.ReqGetLog = Depends()):
|
|
lines = log.buffer[:req.lines] if req.lines > 0 else log.buffer.copy()
|
|
if req.clear:
|
|
log.buffer.clear()
|
|
return lines
|
|
|
|
def post_log(req: models.ReqPostLog):
|
|
if req.json is not None:
|
|
log.info(f'UI {req.message or ""}: {req.json}')
|
|
elif req.message is not None:
|
|
log.info(f'UI: {req.message}')
|
|
elif req.debug is not None:
|
|
log.debug(f'UI: {req.debug}')
|
|
elif req.error is not None:
|
|
log.error(f'UI: {req.error}')
|
|
return Response(status_code=204)
|
|
|
|
def post_shutdown(background_tasks: BackgroundTasks):
|
|
log.info("Server shutdown request received")
|
|
background_tasks.add_task(os._exit, 0)
|
|
return Response(status_code=204)
|
|
|
|
def post_restart(background_tasks: BackgroundTasks):
|
|
log.info("Server restart request received")
|
|
from installer import restart
|
|
background_tasks.add_task(restart)
|
|
return Response(status_code=204)
|
|
|
|
def get_cmd_flags():
|
|
return vars(shared.cmd_opts)
|
|
|
|
def get_history(req: models.ReqHistory = Depends()):
|
|
if req.id is not None and ((isinstance(req.id, str) and len(req.id) > 0) or isinstance(req.id, int)):
|
|
_id = str(req.id) if isinstance(req.id, int) else req.id
|
|
res = [item for item in shared.state.state_history if item['id'] == _id]
|
|
else:
|
|
res = shared.state.state_history
|
|
res = [models.ResHistory(**item) for item in res]
|
|
return res
|
|
|
|
def get_storage(req: models.ReqStorage = Depends()):
|
|
from modules.storage import check_storage
|
|
res = check_storage(folders=req.folder,
|
|
types=req.types.split(',') if req.types else None,
|
|
silent=True,
|
|
)
|
|
res = [models.ResStorage(**loc.dict()) for loc in res]
|
|
return res
|
|
|
|
def get_progress(req: models.ReqProgress = Depends()):
|
|
if shared.state.job_count == 0 and shared.state.sampling_step == 0: # truly idle
|
|
return models.ResProgress(id=shared.state.id, progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
|
|
shared.state.do_set_current_image()
|
|
current_image = None
|
|
if shared.state.current_image and not req.skip_current_image:
|
|
current_image = helpers.encode_pil_to_base64(shared.state.current_image)
|
|
batch_x = max(shared.state.job_no, 0)
|
|
batch_y = max(shared.state.job_count, 1)
|
|
step_x = max(shared.state.sampling_step, 0)
|
|
prev_steps = max(shared.state.sampling_steps, 1)
|
|
while step_x > shared.state.sampling_steps:
|
|
shared.state.sampling_steps += prev_steps
|
|
step_y = max(shared.state.sampling_steps, 1)
|
|
current = step_y * batch_x + step_x
|
|
total = step_y * batch_y
|
|
progress = min((current / total) if current > 0 and total > 0 else 0, 1)
|
|
time_since_start = time.time() - shared.state.time_start
|
|
eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
|
|
# log.trace(f'get_progress: batch {batch_x}/{batch_y} step {step_x}/{step_y} current {current}/{total} time={time_since_start} eta={eta_relative}')
|
|
# log.trace(shared.state)
|
|
res = models.ResProgress(id=shared.state.id, progress=round(progress, 2), eta_relative=round(eta_relative, 2), current_image=current_image, textinfo=shared.state.textinfo, state=shared.state.dict(), )
|
|
return res
|
|
|
|
def get_status():
|
|
return shared.state.status()
|
|
|
|
def post_interrupt():
|
|
shared.state.interrupt()
|
|
return Response(status_code=204)
|
|
|
|
def post_skip():
|
|
shared.state.skip()
|
|
return Response(status_code=204)
|
|
|
|
def get_model_placement():
|
|
import torch
|
|
from modules.modeldata import model_data
|
|
pipe = model_data.sd_model # raw slot: the shared.sd_model property can trigger a model load
|
|
if pipe is None:
|
|
return {}
|
|
components = getattr(pipe, 'components', None) or ({ 'model': pipe } if isinstance(pipe, torch.nn.Module) else {})
|
|
placement = {}
|
|
seen = set()
|
|
for name, component in components.items():
|
|
if not isinstance(component, torch.nn.Module):
|
|
continue
|
|
devmap = {}
|
|
for tensors in (component.parameters(), component.buffers()):
|
|
for t in tensors:
|
|
ptr = 0 if t.is_meta else t.untyped_storage().data_ptr()
|
|
if ptr:
|
|
if ptr in seen:
|
|
continue
|
|
seen.add(ptr)
|
|
devmap[t.device.type] = devmap.get(t.device.type, 0) + t.numel() * t.element_size()
|
|
if devmap:
|
|
placement[name] = devmap
|
|
return placement
|
|
|
|
|
|
def get_memory():
|
|
try:
|
|
import psutil
|
|
process = psutil.Process(os.getpid())
|
|
res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
|
|
ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
|
|
ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
|
|
except Exception as err:
|
|
ram = { 'error': f'{err}' }
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
s = torch.cuda.mem_get_info()
|
|
system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
|
|
s = dict(torch.cuda.memory_stats(shared.device))
|
|
allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
|
|
reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
|
|
active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
|
|
inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
|
|
warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
|
|
cuda = {
|
|
'system': system,
|
|
'active': active,
|
|
'allocated': allocated,
|
|
'reserved': reserved,
|
|
'inactive': inactive,
|
|
'events': warnings,
|
|
}
|
|
else:
|
|
cuda = { 'error': 'unavailable' }
|
|
except Exception as err:
|
|
cuda = { 'error': f'{err}' }
|
|
try:
|
|
model = get_model_placement() # walk can race a model reload or offload rewrap; report the error and keep the endpoint alive
|
|
except Exception as err:
|
|
model = { 'error': f'{err}' }
|
|
return models.ResMemory(ram = ram, cuda = cuda, model = model)
|