From ab13d10e97cfa1c235ba54bd230dc2d9dc7f365c Mon Sep 17 00:00:00 2001 From: Abdullah Alfaraj Date: Sat, 4 Feb 2023 12:56:34 +0300 Subject: [PATCH] add mask expansion and prompt shortcut support to the .ccx package --- .gitignore | 2 + scripts/main.py | 44 +++++++++++-- sdapi_py_re.js | 3 +- server/python_server/img2imgapi.py | 26 ++++++++ server/python_server/serverMain.py | 44 +++++++++++-- utility/sdapi/python_replacement.js | 97 ++++++++++++++++++++++++----- 6 files changed, 188 insertions(+), 28 deletions(-) diff --git a/.gitignore b/.gitignore index cd518ba..2911d27 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,5 @@ experimental/ start_server.sh start_server.bat *.ccx +expanded_mask.png +original_mask.png \ No newline at end of file diff --git a/scripts/main.py b/scripts/main.py index c35b67c..2030913 100644 --- a/scripts/main.py +++ b/scripts/main.py @@ -23,13 +23,14 @@ python_server_full_path = os.path.join(extension_dir,python_server_dir) print("python_server_full_path: ",python_server_full_path) sys.path.insert(0, python_server_full_path) import search +import img2imgapi router = APIRouter() -@router.get("/config") -async def get_state(): - print("hello get /config auto-photoshop-sd") - res = "hello get /config auto-photoshop-sd" - return {"res": res} +# @router.get("/config") +# async def get_state(): +# print("hello get /config auto-photoshop-sd") +# res = "hello get /config auto-photoshop-sd" +# return {"res": res} @router.post('/search/image/') async def searchImage(request:Request): @@ -40,7 +41,7 @@ async def searchImage(request:Request): try: - keywords = json.get('keywords','cute dogs') + keywords = json.get('keywords','cute cats') images = await search.imageSearch(keywords) print(images) @@ -52,6 +53,37 @@ async def searchImage(request:Request): return {"error": "error message: can't preform an image search"} +@router.post('/mask/expansion/') +async def maskExpansionHandler(request:Request): + try: + json = await request.json() + except: + json = {} + + # print("mask expansion json :",json) + try: + # keywords = json.get('keywords','cute dogs') + base64_mask_image = json['mask'] + mask_expansion = json['mask_expansion'] + #convert base64 to img + + await img2imgapi.base64ToPng(base64_mask_image,"original_mask.png")#save a copy of the mask + + mask_image = img2imgapi.b64_2_img(base64_mask_image) + + expanded_mask_img = img2imgapi.maskExpansion(mask_image,mask_expansion) + base64_expanded_mask_image = img2imgapi.img_2_b64(expanded_mask_img) + await img2imgapi.base64ToPng(base64_expanded_mask_image,"expanded_mask.png")#save a copy of the mask + + + return {"mask":base64_expanded_mask_image} + + except: + # print("request",request) + raise Exception(f"couldn't preform mask expansion",json) + # return response + return {"error": "error message: can't preform an mask expansion"} + def on_app_started(demo: gr.Blocks, app: FastAPI): diff --git a/sdapi_py_re.js b/sdapi_py_re.js index ae6a65b..ac07cee 100644 --- a/sdapi_py_re.js +++ b/sdapi_py_re.js @@ -359,7 +359,7 @@ async function savePromptShortcut(prompt_shortcut) { async function setInpaintMaskWeight(value) { const full_url = 'http://127.0.0.1:8000/sdapi/v1/options' try { - payload = { + const payload = { inpainting_mask_weight: value, } await fetch(full_url, { @@ -407,7 +407,6 @@ async function requestGetOptions() { async function imageSearch(keywords) { let json = {} - const full_url = `${g_sd_url}/sdapi/auto-photoshop-sd/search/image/` try { payload = { diff --git a/server/python_server/img2imgapi.py b/server/python_server/img2imgapi.py index 9bde581..12507bf 100644 --- a/server/python_server/img2imgapi.py +++ b/server/python_server/img2imgapi.py @@ -23,6 +23,32 @@ import time import serverHelper import metadata_to_json +def b64_2_img(base64_image): + image = Image.open(io.BytesIO(base64.b64decode(base64_image.split(",",1)[0]))) + return image + +def maskExpansion(mask_img,mask_expansion): + #only if image exist then try to open it + + + # init_img_mask = Image.open(f"{init_img_dir}/{init_img_mask_name}") + + # if(payload['use_sharp_mask'] == False):# use blurry mask + iteration = mask_expansion + mask_img = applyDilation(mask_img,iteration) + + + return mask_img + +async def base64ToPng(base64_image,image_path): + base64_img_bytes = base64_image.encode('utf-8') + with open(image_path, 'wb') as file_to_save: + decoded_image_data = base64.decodebytes(base64_img_bytes) + file_to_save.write(decoded_image_data) + + + + from PIL import Image, ImageFilter def applyDilation(img,iteration=20,max_filter=3): # img = Image.open("test_image_2.png") diff --git a/server/python_server/serverMain.py b/server/python_server/serverMain.py index 40a9934..22fa666 100644 --- a/server/python_server/serverMain.py +++ b/server/python_server/serverMain.py @@ -16,6 +16,8 @@ import search sd_url = os.environ.get('SD_URL', 'http://127.0.0.1:7860') + + async def txt2ImgRequest(payload): # payload = { # "prompt": "cute cat, kitten", @@ -254,11 +256,11 @@ async def sdapi(path: str, request: Request, response: Response): -async def base64ToPng(base64_image,image_path): - base64_img_bytes = base64_image.encode('utf-8') - with open(image_path, 'wb') as file_to_save: - decoded_image_data = base64.decodebytes(base64_img_bytes) - file_to_save.write(decoded_image_data) +# async def base64ToPng(base64_image,image_path): +# base64_img_bytes = base64_image.encode('utf-8') +# with open(image_path, 'wb') as file_to_save: +# decoded_image_data = base64.decodebytes(base64_img_bytes) +# file_to_save.write(decoded_image_data) @app.post('/save/png/') @@ -274,7 +276,7 @@ async def savePng(request:Request): try: folder = './init_images' image_path = f"{folder}/{json['image_name']}" - await base64ToPng(json['base64'],image_path) + await img2imgapi.base64ToPng(json['base64'],image_path) @@ -305,8 +307,38 @@ async def searchImage(request:Request): # print(f'{request}') return {"error": "error message: can't preform an image search"} +@app.post('/mask/expansion/') +async def maskExpansionHandler(request:Request): + try: + json = await request.json() + except: + json = {} + + try: + # keywords = json.get('keywords','cute dogs') + base64_mask_image = json['mask'] + mask_expansion = json['mask_expansion'] + #convert base64 to img + + await img2imgapi.base64ToPng(base64_mask_image,"original_mask.png")#save a copy of the mask for debugging + + mask_image = img2imgapi.b64_2_img(base64_mask_image) + + expanded_mask_img = img2imgapi.maskExpansion(mask_image,mask_expansion) + base64_expanded_mask_image = img2imgapi.img_2_b64(expanded_mask_img) + await img2imgapi.base64ToPng(base64_expanded_mask_image,"expanded_mask.png")#save a copy of the mask of the expanded_mask for debugging + + + return {"mask":base64_expanded_mask_image} + + except: + print("request",request) + raise Exception(f"couldn't preform mask expansion") # return response + return {"error": "error message: can't preform an mask expansion"} + + @app.post('/history/load') async def loadHistory(request: Request): # {'image_paths','metadata_setting'} diff --git a/utility/sdapi/python_replacement.js b/utility/sdapi/python_replacement.js index 4504eae..2b09b4a 100644 --- a/utility/sdapi/python_replacement.js +++ b/utility/sdapi/python_replacement.js @@ -1,5 +1,7 @@ //how to get environment variable in javascript +const { getPromptShortcut } = require('../html_manip') + function newOutputImageName() { const random_id = Math.floor(Math.random() * 100000000000 + 1) // Date.now() doesn't have enough resolution to avoid duplicate const image_name = `output- ${Date.now()}-${random_id}.png` @@ -7,22 +9,47 @@ function newOutputImageName() { return image_name } +function replacePromptsWithShortcuts( + prompt, + negative_prompt, + prompt_shortcut_dic +) { + // const prompt_shortcut_dict = prompt_shortcut.load() + // prompt_shortcut_dict.update(payload["prompt_shortcut_ui_dict"]) + new_prompt = prompt_shortcut.replaceShortcut(prompt, prompt_shortcut_dic) + // # edit negative prompt, replaceShortcut(negative_prompt) + new_negative_prompt = prompt_shortcut.replaceShortcut( + negative_prompt, + prompt_shortcut_dic + ) + return [new_prompt, new_negative_prompt] +} async function txt2ImgRequest(payload) { console.log('payload:', payload) + // if (payload['use_prompt_shortcut']) { + // // const prompt_shortcut_dict = prompt_shortcut.load() + // // prompt_shortcut_dict.update(payload["prompt_shortcut_ui_dict"]) + // payload['prompt'] = prompt_shortcut.replaceShortcut( + // payload['prompt'], + // payload['prompt_shortcut_ui_dict'] + // ) + // // # edit negative prompt, replaceShortcut(negative_prompt) + // payload['negative_prompt'] = prompt_shortcut.replaceShortcut( + // payload['negative_prompt'], + // payload['prompt_shortcut_ui_dict'] + // ) + // } if (payload['use_prompt_shortcut']) { - // const prompt_shortcut_dict = prompt_shortcut.load() - // prompt_shortcut_dict.update(payload["prompt_shortcut_ui_dict"]) - payload['prompt'] = prompt_shortcut.replaceShortcut( + const [new_prompt, new_negative_prompt] = replacePromptsWithShortcuts( payload['prompt'], - payload['prompt_shortcut_ui_dict'] - ) - // # edit negative prompt, replaceShortcut(negative_prompt) - payload['negative_prompt'] = prompt_shortcut.replaceShortcut( payload['negative_prompt'], payload['prompt_shortcut_ui_dict'] ) + payload['prompt'] = new_prompt + payload['negative_prompt'] = new_negative_prompt } + const endpoint = 'sdapi/v1/txt2img' try { console.log('txt2ImgRequest(): about to send a fetch request') @@ -93,6 +120,35 @@ async function txt2ImgRequest(payload) { // const request_path = '/sdapi/v1/txt2img' } +async function maskExpansionRequest(original_mask, mask_expansion_value) { + // const endpoint = 'sdapi/v1/img2img' + // const full_url = `${g_sd_url}/${endpoint}` + try { + const payload = { + mask: original_mask, + mask_expansion: mask_expansion_value, + } + // const full_url = 'http://127.0.0.1:8000/mask/expansion/' + + const full_url = `${g_sd_url}/sdapi/auto-photoshop-sd/mask/expansion/` + let request = await fetch(full_url, { + method: 'POST', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + }, + body: JSON.stringify(payload), + // "body": payload + }) + + let r = await request.json() + + console.log('maskExpansionRequest json:', r) + return r['mask'] + } catch (e) { + console.warn(e) + } +} async function img2ImgRequest(sd_url, payload) { console.log('payload:', payload) @@ -105,7 +161,15 @@ async function img2ImgRequest(sd_url, payload) { // # edit negative prompt, replaceShortcut(negative_prompt) // payload['negative_prompt'] = prompt_shortcut.replaceShortcut(payload['negative_prompt'],prompt_shortcut_dict) // } - + if (payload['use_prompt_shortcut']) { + const [new_prompt, new_negative_prompt] = replacePromptsWithShortcuts( + payload['prompt'], + payload['negative_prompt'], + payload['prompt_shortcut_ui_dict'] + ) + payload['prompt'] = new_prompt + payload['negative_prompt'] = new_negative_prompt + } // init_img_dir = "./init_images" // init_img_name = payload['init_image_name'] // init_img = Image.open(f"{init_img_dir}/{init_img_name}") @@ -115,15 +179,19 @@ async function img2ImgRequest(sd_url, payload) { // init_img_mask_name = payload.get('init_image_mask_name',"") // #only if image exist then try to open it + // if(len(init_img_mask_name) > 0): // init_img_mask = Image.open(f"{init_img_dir}/{init_img_mask_name}") - // if(payload['use_sharp_mask'] == False):# use blurry mask - // iteration = payload['mask_expansion'] - // init_img_mask = applyDilation(init_img_mask,iteration) - - // init_img_mask_str = img_2_b64(init_img_mask) - // payload['mask'] = init_img_mask_str #there is only one mask, unlike 'init_images' which is of type array + if (payload['use_sharp_mask'] === false && payload['mask']) { + //only if mask is available and sharp_mask is off + // use blurry and expanded mask + const iterations = payload['mask_expansion'] + const mask = await maskExpansionRequest(payload['mask'], iterations) + if (mask) { + payload['mask'] = mask + } + } // print(type(init_img_str)) // #request the images to be generated @@ -279,4 +347,5 @@ module.exports = { txt2ImgRequest, img2ImgRequest, loadHistory, + maskExpansionRequest, }