diff --git a/koboldcpp.py b/koboldcpp.py index ce699e7dd..dcdf05324 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -1639,6 +1639,18 @@ def sd_load_model(model_filename,vae_filename,lora_filename,t5xxl_filename,clipl ret = handle.sd_load_model(inputs) return ret +def sd_oai_tranform_params(genparams): + size = genparams.get('size', "512x512") + if size and size!="": + pattern = r'^\D*(\d+)x(\d+)$' + match = re.fullmatch(pattern, size) + if match: + width = int(match.group(1)) + height = int(match.group(2)) + genparams["width"] = width + genparams["height"] = height + return genparams + def sd_comfyui_tranform_params(genparams): promptobj = genparams.get('prompt', None) if promptobj and isinstance(promptobj, dict): @@ -3700,6 +3712,7 @@ Change Mode
api_format = 0 #1=basic,2=kai,3=oai,4=oai-chat,5=interrogate,6=ollama,7=ollamachat is_imggen = False is_comfyui_imggen = False + is_oai_imggen = False is_transcribe = False is_tts = False is_embeddings = False @@ -3785,10 +3798,12 @@ Change Mode
api_format = 6 elif self.path.endswith('/api/chat'): #ollama api_format = 7 - elif self.path=="/prompt" or self.path.endswith('/sdapi/v1/txt2img') or self.path.endswith('/sdapi/v1/img2img'): + elif self.path=="/prompt" or self.path.endswith('/v1/images/generations') or self.path.endswith('/sdapi/v1/txt2img') or self.path.endswith('/sdapi/v1/img2img'): is_imggen = True if self.path=="/prompt": is_comfyui_imggen = True + elif self.path.endswith('/v1/images/generations'): + is_oai_imggen = True elif self.path.endswith('/api/extra/transcribe') or self.path.endswith('/v1/audio/transcriptions'): is_transcribe = True elif self.path.endswith('/api/extra/tts') or self.path.endswith('/v1/audio/speech') or self.path.endswith('/tts_to_audio'): @@ -3898,6 +3913,8 @@ Change Mode
if is_comfyui_imggen: lastgeneratedcomfyimg = b'' genparams = sd_comfyui_tranform_params(genparams) + elif is_oai_imggen: + genparams = sd_oai_tranform_params(genparams) gen = sd_generate(genparams) genresp = None if is_comfyui_imggen: @@ -3906,6 +3923,8 @@ Change Mode
else: lastgeneratedcomfyimg = b'' genresp = (json.dumps({"prompt_id": "12345678-0000-0000-0000-000000000001","number": 0,"node_errors":{}}).encode()) + elif is_oai_imggen: + genresp = (json.dumps({"created":int(time.time()),"data":[{"b64_json":gen}],"background":"opaque","output_format":"png","size":"1024x1024","quality":"medium"}).encode()) else: genresp = (json.dumps({"images":[gen],"parameters":{},"info":""}).encode()) self.send_response(200)