mirror of
https://github.com/vladmandic/automatic
synced 2026-09-17 16:24:33 +02:00
Pablo Hellmann
237 lines
8.2 KiB
Python
237 lines
8.2 KiB
Python
import base64
|
|
import os
|
|
import io
|
|
import time
|
|
import asyncio
|
|
from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
|
|
from starlette.websockets import WebSocket, WebSocketState
|
|
import modules.shared as shared
|
|
from modules.logger import log
|
|
|
|
|
|
current_task = None
|
|
pending_tasks = {}
|
|
finished_tasks = []
|
|
recorded_results = []
|
|
recorded_results_limit = 2
|
|
debug = os.environ.get('SD_PREVIEW_DEBUG', None) is not None
|
|
debug_log = log.trace if debug else lambda *args, **kwargs: None
|
|
|
|
|
|
class PreviewManager:
|
|
def __init__(self):
|
|
self.active: list[WebSocket] = []
|
|
self._loop = None
|
|
|
|
async def connect(self, ws: WebSocket):
|
|
await ws.accept()
|
|
self._loop = asyncio.get_running_loop()
|
|
self.active.append(ws)
|
|
debug_log(f'Preview WS connect: client={ws.client.host} total={len(self.active)}')
|
|
|
|
def disconnect(self, ws: WebSocket):
|
|
try:
|
|
self.active.remove(ws)
|
|
except ValueError:
|
|
pass
|
|
debug_log(f'Preview WS disconnect: client={ws.client.host} total={len(self.active)}')
|
|
|
|
async def broadcast(self, data: dict):
|
|
disconnected = []
|
|
for ws in list(self.active):
|
|
try:
|
|
if ws.client_state == WebSocketState.CONNECTED:
|
|
await ws.send_json(data)
|
|
else:
|
|
disconnected.append(ws)
|
|
except Exception:
|
|
disconnected.append(ws)
|
|
for ws in disconnected:
|
|
try:
|
|
self.active.remove(ws)
|
|
except ValueError:
|
|
pass
|
|
|
|
def push(self, data: dict):
|
|
if self._loop is None or not self.active:
|
|
return
|
|
try:
|
|
asyncio.run_coroutine_threadsafe(self.broadcast(data), self._loop)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
preview_manager = PreviewManager()
|
|
|
|
|
|
def push_progress(step: int, steps: int, progress: float, active: bool = True):
|
|
try:
|
|
data = {
|
|
'type': 'progress',
|
|
'step': step,
|
|
'steps': steps,
|
|
'progress': progress,
|
|
'active': active,
|
|
'job': shared.state.job,
|
|
'paused': shared.state.paused,
|
|
}
|
|
preview_manager.push(data)
|
|
except Exception as e:
|
|
debug_log(f'Progress push error: {e}')
|
|
|
|
|
|
def push_live_preview(image, step: int, steps: int, progress: float):
|
|
try:
|
|
buffered = io.BytesIO()
|
|
image.save(buffered, format='jpeg', quality=60)
|
|
b64 = base64.b64encode(buffered.getvalue()).decode('ascii')
|
|
data = {
|
|
'type': 'preview',
|
|
'live_preview': f'data:image/jpeg;base64,{b64}',
|
|
'id_live_preview': shared.state.id_live_preview,
|
|
'step': step,
|
|
'steps': steps,
|
|
'progress': progress,
|
|
'job': shared.state.job,
|
|
}
|
|
preview_manager.push(data)
|
|
except Exception as e:
|
|
debug_log(f'Preview push error: {e}')
|
|
|
|
|
|
def push_complete():
|
|
try:
|
|
data = {
|
|
'type': 'complete',
|
|
'active': False,
|
|
'job': shared.state.job,
|
|
}
|
|
preview_manager.push(data)
|
|
except Exception as e:
|
|
debug_log(f'Complete push error: {e}')
|
|
|
|
|
|
def start_task(id_task):
|
|
global current_task # pylint: disable=global-statement
|
|
current_task = id_task
|
|
pending_tasks.pop(id_task, None)
|
|
|
|
|
|
def record_results(id_task, res):
|
|
recorded_results.append((id_task, res))
|
|
if len(recorded_results) > recorded_results_limit:
|
|
recorded_results.pop(0)
|
|
|
|
|
|
def finish_task(id_task):
|
|
global current_task # pylint: disable=global-statement
|
|
if current_task == id_task:
|
|
current_task = None
|
|
finished_tasks.append(id_task)
|
|
if len(finished_tasks) > 16:
|
|
finished_tasks.pop(0)
|
|
push_complete()
|
|
|
|
|
|
def add_task_to_queue(id_job):
|
|
pending_tasks[id_job] = time.time()
|
|
|
|
|
|
class ProgressRequest(BaseModel):
|
|
id_task: str = Field(default=None, title="Task ID", description="id of the task to get progress for")
|
|
id_live_preview: int = Field(default=-1, title="Live preview image ID", description="id of last received last preview image")
|
|
|
|
|
|
class InternalProgressResponse(BaseModel):
|
|
job: str = Field(default=None, title="Job name", description="Internal job name")
|
|
textinfo: str|None = Field(default=None, title="Info text", description="Info text used by WebUI.")
|
|
# status fields
|
|
active: bool = Field(title="Whether the task is being worked on right now")
|
|
queued: bool = Field(title="Whether the task is in queue")
|
|
paused: bool = Field(title="Whether the task is paused")
|
|
completed: bool = Field(title="Whether the task has already finished")
|
|
debug: bool = Field(title="Debug logging level")
|
|
# raw fields
|
|
step: int = Field(default=None, title="Current step", description="Current step of the task")
|
|
steps: int = Field(default=None, title="Total steps", description="Total number of steps")
|
|
batch_no: int = Field(default=None, title="Current batch", description="Current batch")
|
|
batch_count: int = Field(default=None, title="Total batches", description="Total number of batches")
|
|
# calculated fields
|
|
progress: float = Field(default=None, title="Progress", description="The progress with a range of 0 to 1")
|
|
eta: float|None = Field(default=None, title="ETA in secs")
|
|
# image fields
|
|
live_preview: str|None = Field(default=None, title="Live preview image", description="Current live preview; a data: uri")
|
|
id_live_preview: int|None = Field(default=None, title="Live preview image ID", description="Send this together with next request to prevent receiving same image")
|
|
|
|
|
|
def api_progress(req: ProgressRequest):
|
|
active = (req.id_task == current_task) or (req.id_task == 'undefined' and len(shared.state.job) > 0)
|
|
queued = req.id_task in pending_tasks
|
|
completed = req.id_task in finished_tasks
|
|
paused = shared.state.paused
|
|
step = max(shared.state.sampling_step, 0)
|
|
steps = max(shared.state.sampling_steps, 1)
|
|
batch_no = max(shared.state.batch_no, 0)
|
|
batch_count = max(shared.state.batch_count, 0)
|
|
|
|
current = step / steps if step > 0 and steps > 0 else 0
|
|
batch = batch_no / batch_count if batch_no > 0 and batch_count > 0 else 1
|
|
progress = round(min(1, current * batch), 2)
|
|
|
|
elapsed = time.time() - shared.state.time_start if shared.state.time_start is not None else 0
|
|
predicted = elapsed / progress if progress > 0 else None
|
|
eta = predicted - elapsed if predicted is not None else None
|
|
id_live_preview = req.id_live_preview
|
|
live_preview = None
|
|
textinfo = shared.state.textinfo
|
|
if not active:
|
|
id_live_preview = -1
|
|
textinfo = "Queued..." if queued else "Waiting..."
|
|
|
|
debug_log(f'Progress: job="{shared.state.job}" active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} job={shared.state.preview_job} elapsed={elapsed:.3f}')
|
|
|
|
if active:
|
|
shared.state.set_current_image()
|
|
if shared.state.current_image is not None and req.id_live_preview != shared.state.id_live_preview:
|
|
live_preview = shared.state.live_preview_b64
|
|
|
|
id_live_preview = shared.state.id_live_preview
|
|
|
|
res = InternalProgressResponse(
|
|
job=shared.state.job,
|
|
textinfo=textinfo,
|
|
active=active,
|
|
queued=queued,
|
|
paused=paused,
|
|
completed=completed,
|
|
debug=debug,
|
|
progress=progress,
|
|
step=step,
|
|
steps=steps,
|
|
batch_no=batch_no,
|
|
batch_count=batch_count,
|
|
job_timestamp=shared.state.time_start,
|
|
eta=eta,
|
|
live_preview=live_preview,
|
|
id_live_preview=id_live_preview,
|
|
)
|
|
return res
|
|
|
|
|
|
def setup_progress_api():
|
|
shared.api.add_api_route("/internal/progress", api_progress, methods=["POST"], response_model=InternalProgressResponse)
|
|
|
|
@shared.api.app.websocket("/ws/preview")
|
|
async def ws_preview(ws: WebSocket):
|
|
await preview_manager.connect(ws)
|
|
try:
|
|
while True:
|
|
data = await ws.receive_text()
|
|
if data == 'ping':
|
|
await ws.send_text('pong')
|
|
elif data == 'end':
|
|
break
|
|
except Exception:
|
|
pass
|
|
preview_manager.disconnect(ws)
|