File size: 2,894 Bytes
0163a2c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
import io
import sys
import fastapi
import gradio as gr
from pydantic import BaseModel, Field
import scripts.shared as shared
from scripts.utilities import run_python
proc = None
outputs = []
def alive():
return proc is not None
def initialize_runner(script_file, tmpls, opts):
run_button = gr.Button(
"Run",
variant="primary",
elem_id=f"kohya_sd_webui__{shared.current_tab}_run_button",
)
stop_button = gr.Button(
"Stop",
variant="secondary",
elem_id=f"kohya_sd_webui__{shared.current_tab}_stop_button",
)
get_templates = lambda: tmpls() if callable(tmpls) else tmpls
get_options = lambda: opts() if callable(opts) else opts
def run(args):
global proc
global outputs
if alive():
return
proc = run_python(script_file, get_templates(), get_options(), args)
reader = io.TextIOWrapper(proc.stdout, encoding="utf-8-sig")
line = ""
while proc is not None and proc.poll() is None:
try:
char = reader.read(1)
if shared.cmd_opts.enable_console_log:
sys.stdout.write(char)
if char == "\n":
outputs.append(line)
line = ""
continue
line += char
except:
()
proc = None
def stop():
global proc
print("killed the running process")
proc.kill()
proc = None
def init():
run_button.click(
run,
set(get_options().values()),
)
stop_button.click(stop)
return init
class GetOutputRequest(BaseModel):
output_index: int = Field(
default=0, title="Index of the beginning of the log to retrieve"
)
clear_terminal: bool = Field(
default=False, title="Whether to clear the terminal"
)
class GetOutputResponse(BaseModel):
outputs: list = Field(title="List of terminal output")
class ProcessAliveResponse(BaseModel):
alive: bool = Field(title="Whether the process is running.")
def api_get_outputs(req: GetOutputRequest):
i = req.output_index
if req.clear_terminal:
global outputs
outputs = []
out = outputs[i:] if len(outputs) > i else []
return GetOutputResponse(outputs=out)
def api_get_isalive(req: fastapi.Request):
return ProcessAliveResponse(alive=alive())
def initialize_api(app: fastapi.FastAPI):
app.add_api_route(
"/internal/extensions/kohya-sd-scripts-webui/terminal/outputs",
api_get_outputs,
methods=["POST"],
response_model=GetOutputResponse,
)
app.add_api_route(
"/internal/extensions/kohya-sd-scripts-webui/process/alive",
api_get_isalive,
methods=["GET"],
response_model=ProcessAliveResponse,
)
|