diff --git a/expose.cpp b/expose.cpp index b1befd572..5ca4a7813 100644 --- a/expose.cpp +++ b/expose.cpp @@ -217,6 +217,10 @@ extern "C" { sdtype_abort_generation(); } + sd_info_outputs sd_get_ongoing_generation_info() + { + return sdtype_get_ongoing_generation_info(); + } bool whisper_load_model(const whisper_load_model_inputs inputs) { diff --git a/koboldcpp.py b/koboldcpp.py index 0a71a8d51..604aabe88 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -36,7 +36,7 @@ import urllib.parse import urllib.request from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timezone -from typing import Tuple +from typing import Any, Dict, Optional, Tuple import shutil import subprocess import gzip @@ -100,6 +100,7 @@ mmprojName = None lastgeneratedcomfyimg = b'' lastgeneratedcachedimg = b'' lastgeneratedcachedimgkey = b'' +currgenimgkey = '' lastuploadedcomfyimg = b'' fullsdmodelpath = "" #if empty, it's not initialized password = "" #if empty, no auth key required @@ -1000,6 +1001,8 @@ def init_library(): handle.sd_get_info.restype = sd_info_outputs handle.sd_abort_generation.argtypes = [] handle.sd_abort_generation.restype = None + handle.sd_get_ongoing_generation_info.argtypes = [] + handle.sd_get_ongoing_generation_info.restype = sd_info_outputs handle.whisper_load_model.argtypes = [whisper_load_model_inputs] handle.whisper_load_model.restype = ctypes.c_bool handle.whisper_generate.argtypes = [whisper_generation_inputs] @@ -2970,6 +2973,102 @@ def sd_generate(genparams): info["job_timestamp"] = job_timestamp return {"animated": animated, "data":data_main, "data_extra":data_extra, "final_frame":final_frame, "info": info} +def sd_get_ongoing_generation_info(): + info = handle.sd_get_ongoing_generation_info() + if info.status == 0: + try: + return json.loads(info.data) + except Exception: + print("An error occurred while decoding sd ongoig generation info") + else: + print("An error occurred while getting sd ongoig generation info") + return {} + +def build_a1111_progress_response( + status: str, + step_count: int = 0, + total_steps: int = 0, + elapsed_time: float = 0.0, + current_image: Optional[str] = None, +) -> Dict[str, Any]: + """ + Maps generation state from the backend to the AUTOMATIC1111 ProgressResponse format. + + :param status: one of 'idle', 'conditioning', 'diffusing', 'vae' + :param step_count: current sampling step + :param total_steps: total sampling steps for the job + :param elapsed_time: seconds since the generation started + :param current_image: base64 encoded string of the preview image, or None + """ + + # "idle" state signature from A1111 + progress = 0.0 + eta_relative = 0.0 + state = { + "job_count": 0, + "job_no": 0, + "sampling_step": 0, + "sampling_steps": 0, + "interrupted": False, + "skipped": False + } + textinfo = None + + if status != 'idle': + + # single job (queue is client-side) + state["job_count"] = 1 + state["job_no"] = 0 + state["sampling_step"] = step_count + state["sampling_steps"] = total_steps + + # calculate progress estimate: fixed portion allocated to conditioning, + # and decoding estimate based on the total number of steps + conditioning_progress = 0.03 + decoding_steps = 0.8 + effective_total = total_steps + decoding_steps + + progress_per_step = (1.0 - conditioning_progress) / effective_total if effective_total > 0 else 0.0 + + if status == 'conditioning': + # A1111 sets progress to > 0.0 during initialization to show activity + # and avoid dividing by zero in the ETA calculation later. + progress = conditioning_progress + textinfo = "Status: conditioning" + elif status == 'diffusing': + progress = conditioning_progress + (step_count * progress_per_step) + textinfo = f"Status: diffusing, Step: {step_count}/{total_steps}" + elif status == 'decoding': + # A1111 sets this as a fixed 0.99, but we keep it proportional + progress = conditioning_progress + (total_steps * progress_per_step) + textinfo = "Status: decoding" + else: + # shouldn't happen + progress = 0.99 + + # note A1111 caps progress below 1.0 until the final image is actually + # returned by the blocking txt2img endpoint, preventing the progress bar + # from jumping to 100% early + + eta_relative = (elapsed_time / progress) - elapsed_time + + return { + "progress": round(progress, 4), + "eta_relative": round(eta_relative, 2), + "state": state, + "current_image": current_image, + "textinfo": textinfo + } + +def a1111_progress_response(preview=False): + status = sd_get_ongoing_generation_info() + result = build_a1111_progress_response( + status.get('status', 0), + status.get('step', 0), + status.get('steps', 1), + status.get('step_time', 1), + preview and status.get('preview') or None) + return result def whisper_load_model(model_filename): global args @@ -6221,7 +6320,7 @@ Change Mode
def do_GET(self): global embedded_kailite, embedded_kcpp_docs, embedded_kcpp_sdui, embedded_kailite_gz, embedded_kcpp_docs_gz, embedded_kcpp_sdui_gz, embedded_lcpp_ui_gz, embedded_musicui, embedded_musicui_gz global last_req_time, start_time, cached_chat_template, cached_sd_info, has_vision_support, has_audio_support, has_whisper, friendlymodelname - global savedata_obj, has_multiplayer, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, maxctx, maxhordelen, friendlymodelname, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey, friendlysdmodelname, fullsdmodelpath, password, friendlyembeddingsmodelname, voicelist + global savedata_obj, has_multiplayer, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, maxctx, maxhordelen, friendlymodelname, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, currgenimgkey, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey, friendlysdmodelname, fullsdmodelpath, password, friendlyembeddingsmodelname, voicelist global autoswapmode, textName, sttName, ttsName, embedName, musicName, imageName, mmprojName clean_path = self.path.split("?")[0] #for cases where we do not want query params @@ -6481,6 +6580,15 @@ Change Mode
response_body = lastgeneratedcachedimg else: response_body = None + elif clean_path.startswith('/sdapi/v1/progress'): + parsed_url = urllib.parse.urlparse(self.path) + parsed_dict = urllib.parse.parse_qs(parsed_url.query) + genkey = parsed_dict.get('genkey', [''])[0] + skip_current_image = bool(parsed_dict.get('skip_current_image', False)) + # with no auth, reveal status without preview image + auth = bool(genkey and genkey==currgenimgkey) + info = a1111_progress_response(auth and not skip_current_image) + response_body = json.dumps(info).encode() elif clean_path=='/history' or clean_path=='/api/history' or clean_path.startswith('/api/history/') or clean_path.startswith('/history/'): #emulate comfyui modelNameToReturn = friendlysdmodelname if autoswapmode and imageName is not None: @@ -6616,7 +6724,7 @@ Change Mode
def do_POST(self): global thinkformats - global modelbusy, batched_request_runner_count, requestsinqueue, currentusergenkey, totalgens, pendingabortkey, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, net_save_slots, has_vision_support, savestate_limit, mcp_lock + global modelbusy, batched_request_runner_count, requestsinqueue, currentusergenkey, totalgens, pendingabortkey, lastuploadedcomfyimg, lastgeneratedcomfyimg, lastgeneratedcachedimg, lastgeneratedcachedimgkey, currgenimgkey, multiplayer_turn_major, multiplayer_turn_minor, multiplayer_story_data_compressed, multiplayer_dataformat, multiplayer_lastactive, net_save_slots, has_vision_support, savestate_limit, mcp_lock global autoswapmode, textName, sttName, ttsName, embedName, musicName, imageName, mmprojName contlenstr = self.headers['content-length'] content_length = 0 @@ -7579,6 +7687,7 @@ Change Mode
try: lastgeneratedcachedimg = b'' lastgeneratedcachedimgkey = '' + currgenimgkey = genparams.get('genkey', '') if is_comfyui_imggen: lastgeneratedcomfyimg = b'' genparams = sd_comfyui_tranform_params(genparams) @@ -7600,6 +7709,7 @@ Change Mode
gendatextra = gen["data_extra"] genfinalframe = gen["final_frame"] geninfo = json.dumps(gen["info"]) # sdapi really expects a stringified JSON + currgenimgkey = '' genresp = None if gendat: lastgeneratedcachedimg = base64.b64decode(gendat) diff --git a/model_adapter.h b/model_adapter.h index d33da9869..d4b144533 100644 --- a/model_adapter.h +++ b/model_adapter.h @@ -112,6 +112,7 @@ bool sdtype_load_model(const sd_load_model_inputs inputs); sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs); sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs); sd_info_outputs sdtype_get_info(); +sd_info_outputs sdtype_get_ongoing_generation_info(); void sdtype_abort_generation(); bool whispertype_load_model(const whisper_load_model_inputs inputs); diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 04d5af70f..d0b57f14e 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -2,11 +2,13 @@ #include #include #include +#include #include #include #include #include #include +#include #include #include @@ -158,6 +160,20 @@ static bool photomaker_enabled = false; static bool is_vid_model = false; static bool remove_limits = false; +struct gendata_st { + int status; + int step; + double step_time; + std::string preview; +}; + +struct { + std::mutex mux; + std::chrono::steady_clock::time_point start_time; + int steps; + gendata_st gendata; +} geninfo; + static struct { std::string data; std::string data_extra; @@ -273,12 +289,29 @@ std::string load_gpt_oss_vocab_json() return load_embd_file(cache, "embd_res/gpt_oss_vocab_json.embd"); } +static void progress_callback(int step, int steps, float time, void* data) +{ + (void) data; + if(sd_is_quiet) return; + const char* unit = "s/it"; + float speed = time; + if (speed < 1.0f && speed > 0.f) { + speed = 1.0f / speed; + unit = "it/s"; + } + printf("Generating image: %d/%d steps, %.2f%s\n", step, steps, speed, unit); + fflush(stdout); +} + +static void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data); + static bool is_video_model(kcpp_sd::model_info info) { return info.is_wan || info.is_ltx || info.is_minimaxh3; } bool sdtype_load_model(const sd_load_model_inputs inputs) { + sd_is_quiet = inputs.quiet; set_sd_quiet(sd_is_quiet); executable_path = sd_get_u8path(inputs.executable_path); @@ -541,6 +574,13 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { } } + sd_set_preview_callback(step_callback, PREVIEW_PROJ, 1, true, false, nullptr); + + if (sddebugmode) { + // the default progress bar would become intermingled with the debug log + sd_set_progress_callback(progress_callback, nullptr); + } + return true; } @@ -986,10 +1026,25 @@ void sdtype_abort_generation() { sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) { + struct CleanupInfoOnExit { + ~CleanupInfoOnExit() { + std::lock_guard lock(geninfo.mux); + geninfo.gendata.status = 0; + } + } cleanup_info_on_exit; + if(sd_ctx == nullptr || sd_params == nullptr) { return sd_generation.error("Warning: KCPP image generation not initialized!"); } + + { + std::lock_guard lock(geninfo.mux); + geninfo.start_time = std::chrono::steady_clock::now(); + geninfo.gendata.status = 1; + geninfo.steps = inputs.sample_steps; + } + sd_image_t * results = nullptr; int generated_num_results = 0; @@ -1341,6 +1396,11 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) sd_audio_t* generated_audio = nullptr; sd_audio_t input_audio = {0, 0, 0, nullptr}; + { + std::lock_guard lock(geninfo.mux); + geninfo.steps = inputs.sample_steps; + } + if(is_vid_model) { std::vector control_frames; //empty for now @@ -1663,6 +1723,32 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) return sd_generation.outputs(1); } +static inline double get_time_delta(const std::chrono::steady_clock::time_point& start) { + auto now = std::chrono::steady_clock::now(); + return std::chrono::duration(now - start).count(); +} + +static void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) +{ + gendata_st gendata; + gendata.status = 2; + if (frame_count == 1) { + gendata.preview = raw_image_to_png_base64(*image); + } else { + uint8_t * out_data = nullptr; + size_t out_len = 0; + create_gif_buf_from_sd_images_msf(image, frame_count, 16, &out_data,&out_len); + gendata.preview = kcpp_base64_encode(out_data, out_len); + } + gendata.step = step; + gendata.step_time = get_time_delta(geninfo.start_time); + + std::lock_guard lock(geninfo.mux); + if (step == geninfo.steps) + gendata.status = 3; + geninfo.gendata = gendata; +} + sd_generation_outputs sdtype_upscale(const sd_upscale_inputs inputs) { sd_generation.reset(); @@ -1755,3 +1841,38 @@ sd_info_outputs sdtype_get_info() return output; } +sd_info_outputs sdtype_get_ongoing_generation_info() +{ + double elapsed_time; + int steps; + gendata_st gendata; + { + std::lock_guard lock(geninfo.mux); + gendata = geninfo.gendata; + elapsed_time = get_time_delta(geninfo.start_time); + steps = geninfo.steps; + } + + nlohmann::json j; + j["steps"] = steps; + j["elapsed_time"] = elapsed_time; + j["step"] = gendata.step; + j["step_time"] = gendata.step_time; + if (gendata.status == 1) + j["status"] = "conditioning"; + else if (gendata.status == 2) + j["status"] = "diffusing"; + else if (gendata.status == 3) + j["status"] = "decoding"; + else + j["status"] = "idle"; + j["preview"] = gendata.preview; + + static std::string recent_info; + recent_info = j.dump(); + sd_info_outputs output; + output.status = 0; + output.data = recent_info.c_str(); + return output; +} +