diff --git a/javascript/descriptions.js b/javascript/descriptions.js index 5e29885..57eb6e0 100644 --- a/javascript/descriptions.js +++ b/javascript/descriptions.js @@ -11,34 +11,38 @@ onUiUpdate(() => { if (!app || app === document) return; const descs = { - '#dumpunet-features-checkbox': 'Extract U-Net features and add their maps to output images.', - '#dumpunet-features-layer': 'U-Net layers (IN00-IN11, M00, OUT00-OUT11) which features should be extracted. See tooltip for notations.', - '#dumpunet-features-steps': 'Steps which U-Net features should be extracted. See tooltip for notations', - '#dumpunet-features-dumppath': 'Raw binary files are dumped to here, one image per step per layer.', + '#dumpunet-{}-features-checkbox': 'Extract U-Net features and add their maps to output images.', + '#dumpunet-{}-features-layer': 'U-Net layers (IN00-IN11, M00, OUT00-OUT11) which features should be extracted. See tooltip for notations.', + '#dumpunet-{}-features-steps': 'Steps which U-Net features should be extracted. See tooltip for notations', + '#dumpunet-{}-features-dumppath': 'Raw binary files are dumped to here, one image per step per layer.', - '#dumpunet-layerprompt-checkbox': 'When checked, (~: ... :~) notation is enabled.', - '#dumpunet-layerprompt-stdout-checkbox': 'When checked, print prompts for each batch to stdout.', + '#dumpunet-{}-layerprompt-checkbox': 'When checked, (~: ... :~) notation is enabled.', + '#dumpunet-{}-layerprompt-stdout-checkbox': 'When checked, print prompts for each batch to stdout.', }; const hints = { - '#dumpunet-features-layer textarea': 'IN00: add one layer to output\nIN00,IN01: add layers to output\nIN00-IN02: add range to output\nIN00-OUT05(+2): add range to output with specified steps\n', - '#dumpunet-features-steps textarea': '5: extracted at steps=5\n5,10: extracted at steps=5 and steps=10\n5-10: extracted when step is in 5..10 (inclusive)\n5-10(+2): extracts when step is 5,7,9\n', + '#dumpunet-{}-features-layer textarea': 'IN00: add one layer to output\nIN00,IN01: add layers to output\nIN00-IN02: add range to output\nIN00-OUT05(+2): add range to output with specified steps\n', + '#dumpunet-{}-features-steps textarea': '5: extracted at steps=5\n5,10: extracted at steps=5 and steps=10\n5-10: extracted when step is in 5..10 (inclusive)\n5-10(+2): extracts when step is 5,7,9\n', }; for (let [k, v] of Object.entries(descs)) { const cont = document.createElement('div'); cont.innerHTML = v; cont.classList.add('dumpunet-description'); - app.querySelector(k).append(cont); + for (let x of ['txt2img', 'img2img']) { + app.querySelector(k.replace('{}', x)).append(cont.cloneNode(true)); + } } for (let [k, v] of Object.entries(hints)) { const cont = document.createElement('pre'); cont.innerHTML = v; cont.classList.add('dumpunet-tooltip'); - const parent = app.querySelector(k).parentNode; - parent.classList.add('dumpunet-tooltip-parent'); - parent.append(cont); + for (let x of ['txt2img', 'img2img']) { + const parent = app.querySelector(k.replace('{}', x)).parentNode; + parent.classList.add('dumpunet-tooltip-parent'); + parent.append(cont.cloneNode(true)); + } } DumpUnet.addDescriptionCallbackCalled = true; diff --git a/javascript/layerinfo.js b/javascript/layerinfo.js index a67f8f9..33f25b1 100644 --- a/javascript/layerinfo.js +++ b/javascript/layerinfo.js @@ -10,36 +10,25 @@ onUiUpdate(() => { const app = gradioApp(); if (!app || app === document) return; - const labels = Array.of(...app.querySelectorAll('#tab_txt2img label')); - const width_label = labels.find(x => x.textContent.trim() === "Width"); - const height_label = labels.find(x => x.textContent.trim() === "Height"); - const steps_label = labels.find(x => x.textContent.trim() === "Sampling Steps"); - if (!width_label || !height_label || !steps_label) return; + const sliders = {}; + for (let mode of ['txt2img', 'img2img']) { + const labels = Array.of(...app.querySelectorAll(`#tab_${mode} label`)); + const width_label = labels.find(x => x.textContent.trim() === "Width"); + const height_label = labels.find(x => x.textContent.trim() === "Height"); + const steps_label = labels.find(x => x.textContent.toLowerCase().trim() === "sampling steps"); + if (!width_label || !height_label || !steps_label) return; - const width_slider = app.querySelector(`#${width_label.htmlFor}`); - const height_slider = app.querySelector(`#${height_label.htmlFor}`); - const steps_slider = app.querySelector(`#${steps_label.htmlFor}`) - if (!width_slider || !height_slider || !steps_slider) return; - //if (+width_slider.dataset.dumpunetHooked && +height_slider.dataset.dumpunetHooked) return - // - //const value_hook = ele => { - // const proto = Object.getPrototypeOf(ele); - // const old_desc = Object.getOwnPropertyDescriptor(proto, 'value'); - // Object.defineProperty(ele, 'value', { - // get: function () { return old_desc.get.apply(this, arguments); }, - // set: function () { - // const old_value = this.value; - // old_desc.set.apply(this, arguments); - // const new_value = this.value; - // const ev = new CustomEvent('imagesizesliderchange', { detail: { old_value: old_value }, bubbles: true }); - // ele.dispatchEvent(ev); - // } - // }); - // ele.dataset.dumpunetHooked = 1; - //}; - // - //value_hook(width_slider); - //value_hook(height_slider); + const width_slider = app.querySelector(`#${width_label.htmlFor}`); + const height_slider = app.querySelector(`#${height_label.htmlFor}`); + const steps_slider = app.querySelector(`#${steps_label.htmlFor}`) + if (!width_slider || !height_slider || !steps_slider) return; + + sliders[mode] = { + width_slider, + height_slider, + steps_slider, + } + } globalThis.DumpUnet.applySizeCallbackCalled = true; @@ -89,51 +78,51 @@ onUiUpdate(() => { return [...new Set(layers)].sort().map(n => layer_names[n]); }; - const layer_input_ele = app.querySelector('#dumpunet-features-layer textarea'); - layer_input_ele.addEventListener('input', e => { - const input = layer_input_ele.value; - const layers = parse_layers(input); - layer_input_ele.style.backgroundColor = layers.length == 0 ? 'pink' : 'white'; - }, false); + for (let mode of ['txt2img', 'img2img']) { + const layer_input_ele = app.querySelector(`#dumpunet-${mode}-features-layer textarea`); + layer_input_ele.addEventListener('input', e => { + const input = layer_input_ele.value; + const layers = parse_layers(input); + layer_input_ele.style.backgroundColor = layers.length == 0 ? 'pink' : 'white'; + }, false); - const update_info = () => { - const layer_input = layer_input_ele.value; - const layers = parse_layers(layer_input); + const update_info = () => { + const layer_input = layer_input_ele.value; + const layers = parse_layers(layer_input); - const - w = +width_slider.value, - h = +height_slider.value, - iw = Math.max(1, Math.ceil(w / 64)), - ih = Math.max(1, Math.ceil(h / 64)); + const + width_slider = sliders[mode].width_slider, + height_slider = sliders[mode].height_slider, + w = +width_slider.value, + h = +height_slider.value, + iw = Math.max(1, Math.ceil(w / 64)), + ih = Math.max(1, Math.ceil(h / 64)); - let html = '[Selected Layer Info]
'; + let html = '[Selected Layer Info]
'; - for (let layer of layers) { - const info = JSON.parse(app.querySelector('#dumpunet-features-layer_setting').textContent)[layer]; - info[0][1] *= ih; - info[0][2] *= iw; - info[1][1] *= ih; - info[1][2] *= iw; - html += ` + for (let layer of layers) { + const info = JSON.parse(app.querySelector(`#dumpunet-${mode}-features-layer_setting`).textContent)[layer]; + info[0][1] *= ih; + info[0][2] *= iw; + info[1][1] *= ih; + info[1][2] *= iw; + html += ` Name:   ${layer}
Input:  (${info[0].join(',')})
Outout: (${info[1].join(',')})
---
`.trim(); - } + } - app.querySelector('#dumpunet-features-layerinfo').innerHTML = html.trim(); - }; + app.querySelector(`#dumpunet-${mode}-features-layerinfo`).innerHTML = html.trim(); + }; - //app.addEventListener('imagesizesliderchange', e => { - // //console.log(e.detail.old_value, e.target.value); - // update_info(); - //}, false); - app.addEventListener('input', update_info, false); - app.addEventListener('change', update_info, false); + app.addEventListener('input', update_info, false); + app.addEventListener('change', update_info, false); - update_info(); + update_info(); + } }; onUiUpdate(DumpUnet.applySizeCallback); diff --git a/javascript/ui.js b/javascript/ui.js index 8edeef0..94cefd2 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -10,9 +10,13 @@ onUiUpdate(() => { const app = gradioApp(); if (!app || app === document) return; - const unet_tab = app.querySelector('#dumpunet-features-tab'), - prompt_tab = app.querySelector('#dumpunet-layerprompt-tab'); - if (!unet_tab || !prompt_tab) return; + const tabs = {}; + for (let mode of ['txt2img', 'img2img']) { + const unet_tab = app.querySelector(`#dumpunet-${mode}-features-tab`), + prompt_tab = app.querySelector(`#dumpunet-${mode}-layerprompt-tab`); + if (!unet_tab || !prompt_tab) return; + tabs[mode] = { unet: unet_tab, prompt: prompt_tab }; + } const disableChildren = function (ele, excepts) { for (let e of ele.querySelectorAll('textarea, input, select')) { @@ -39,14 +43,12 @@ onUiUpdate(() => { }; }; - //unet = apply(unet_tab, '#dumpunet-features-checkbox input[type=checkbox]'); - prompt = apply(prompt_tab, '#dumpunet-layerprompt-checkbox input[type=checkbox]'); - - //unet_tab.addEventListener('change', unet, false); - prompt_tab.addEventListener('change', prompt, false); - - //unet(); - prompt(); + for (let mode of ['txt2img', 'img2img']) { + const tab = tabs[mode]; + const prompt = apply(tab.prompt, `#dumpunet-${mode}-layerprompt-checkbox input[type=checkbox]`); + tab.prompt.addEventListener('change', prompt, false); + prompt(); + } DumpUnet.featureEnablerCallbackCalled = true; }; diff --git a/scripts/dumpunet.py b/scripts/dumpunet.py index cb0feea..4086aaa 100644 --- a/scripts/dumpunet.py +++ b/scripts/dumpunet.py @@ -31,87 +31,90 @@ class Script(scripts.Script): def ui(self, is_img2img): - with gr.Group(elem_id="dumpunet-ui"): - with gr.Tab("U-Net features", elem_id="dumpunet-features-tab"): + ID_PREFIX = "dumpunet" + def id(x): + return f"{ID_PREFIX}-{'img2img' if is_img2img else 'txt2img'}-{x}" + + with gr.Group(elem_id=id("ui")): + with gr.Tab("U-Net features", elem_id=id("features-tab")): unet_features_enabled = gr.Checkbox( label="Extract U-Net features", value=False, - elem_id="dumpunet-features-checkbox" + elem_id=id("features-checkbox") ) - with gr.Group(elem_id="dumpunet-features-image"): + with gr.Group(elem_id=id("features-image")): layer = gr.Textbox( label="Layers", placeholder="eg. IN00-OUT03(+2),OUT10", - elem_id="dumpunet-features-layer", + elem_id=id("features-layer"), ) layer_setting_hidden = gr.HTML( json.dumps(layerinfo.Settings), visible=False, - elem_id="dumpunet-features-layer_setting" + elem_id=id("features-layer_setting") ) steps = gr.Textbox( label="Image saving steps", placeholder="eg. 1,5-20(+5)", - elem_id="dumpunet-features-steps" + elem_id=id("features-steps") ) color = gr.Checkbox( False, label="Use red/blue color map (red=POSITIVE, black=ZERO, blue=NEGATIVE)", - elem_id="dumpunet-features-color" + elem_id=id("features-color") ) - with gr.Group(elem_id="dumpunet-features-tensor"): + with gr.Group(elem_id=id("features-tensor")): path_on = gr.Checkbox( False, label="Dump feature tensors to files", - elem_id="dumpunet-features-dump" + elem_id=id("features-dump") ) path = gr.Textbox( label="Output path", placeholder="eg. /home/hnmr/unet/", - elem_id="dumpunet-features-dumppath" + elem_id=id("features-dumppath") ) with gr.Accordion("Selected Layer Info", open=False): - layer_info = gr.HTML(elem_id="dumpunet-features-layerinfo") + layer_info = gr.HTML(elem_id=id("features-layerinfo")) - with gr.Tab("Layer Prompt", elem_id="dumpunet-layerprompt-tab"): - #with gr.Group(elem_id="dumpunet-layerprompt"): + with gr.Tab("Layer Prompt", elem_id=id("layerprompt-tab")): layerprompt_enabled = gr.Checkbox( label="Enable Layer Prompt", value=False, - elem_id="dumpunet-layerprompt-checkbox" + elem_id=id("layerprompt-checkbox") ) - with gr.Group(elem_id="dumpunet-layerprompt-stdout"): + with gr.Group(elem_id=id("layerprompt-stdout")): layerprompt_show_prompts = gr.Checkbox( label="Show prompts in stdout", value=False, - elem_id="dumpunet-layerprompt-stdout-checkbox" + elem_id=id("layerprompt-stdout-checkbox") ) - with gr.Group(elem_id="dumpunet-layerprompt-diff"): + with gr.Group(elem_id=id("layerprompt-diff")): layerprompt_diff_enabled = gr.Checkbox( label="Output difference map of U-Net features between with and without Layer Prompt", value=False, - elem_id="dumpunet-layerprompt-diff-checkbox" + elem_id=id("layerprompt-diff-checkbox") ) diff_path_on = gr.Checkbox( False, label="Dump difference tensors to files", - elem_id="dumpunet-layerprompt-diff-dump" + elem_id=id("layerprompt-diff-dump") ) diff_path = gr.Textbox( label="Output path", placeholder="eg. /home/hnmr/unet/", - elem_id="dumpunet-layerprompt-diff-dumppath" + elem_id=id("layerprompt-diff-dumppath") ) #with gr.Accordion("Prompt Errors", open=False): @@ -335,9 +338,19 @@ def copy_processing(p: StableDiffusionProcessing): t2i_args = { "enable_hr": p.enable_hr, "denoising_strength": p.denoising_strength, - "firstphase_width": p.firstphase_width, - "firstphase_height": p.firstphase_height + "hr_scale": p.hr_scale, + "hr_upscaler": p.hr_upscaler, + "hr_second_pass_steps": p.hr_second_pass_steps, + "hr_resize_x": p.hr_resize_x, + "hr_resize_y": p.hr_resize_y, } + if p.hr_upscale_to_x != 0 or p.hr_upscale_to_y != 0: + t2i_args.update({ + "firstphase_width": p.width, + "firstphase_height": p.height, + "width": p.hr_upscale_to_x, + "height": p.hr_upscale_to_y, + }) if isinstance(p, StableDiffusionProcessingImg2Img): i2i_args = { @@ -359,8 +372,8 @@ def copy_processing(p: StableDiffusionProcessing): pp = type(p)(**args) pp.prompt_for_display = p.prompt_for_display - pp.paste_to = p.paste_to - pp.color_corrections = p.color_corrections + pp.paste_to = p.paste_to # type: ignore + pp.color_corrections = p.color_corrections # type: ignore pp.sampler_noise_scheduler_override = p.sampler_noise_scheduler_override pp.is_using_inpainting_conditioning = p.is_using_inpainting_conditioning pp.scripts = p.scripts @@ -373,7 +386,8 @@ def copy_processing(p: StableDiffusionProcessing): for attr in [ "sampler", "truncate_x", "truncate_y", - "init_latent", "latent_mask", "mask_for_overlay", "mask", "nmask", "image_conditioning" + "init_latent", "latent_mask", "mask_for_overlay", "mask", "nmask", "image_conditioning", + ]: if hasattr(p, attr): setattr(pp, attr, getattr(p, attr)) diff --git a/style.css b/style.css index dea8c44..0f8723d 100644 --- a/style.css +++ b/style.css @@ -1,16 +1,22 @@ -#dumpunet-features-layerinfo, -#dumpunet-layerprompt-errors { +#dumpunet-txt2img-features-layerinfo, +#dumpunet-img2img-features-layerinfo, +#dumpunet-txt2img-layerprompt-errors, +#dumpunet-img2img-layerprompt-errors { font-family: monospace; } -#dumpunet-features-checkbox, -#dumpunet-layerprompt-checkbox { +#dumpunet-txt2img-features-checkbox, +#dumpunet-img2img-features-checkbox, +#dumpunet-txt2img-layerprompt-checkbox, +#dumpunet-img2img-layerprompt-checkbox { /* background-color: #fff8f0; */ background-color: rgb(255 248 240 / var(--tw-bg-opacity)); } -.dark #dumpunet-features-checkbox, -.dark #dumpunet-layerprompt-checkbox { +.dark #dumpunet-txt2img-features-checkbox, +.dark #dumpunet-img2img-features-checkbox, +.dark #dumpunet-txt2img-layerprompt-checkbox, +.dark #dumpunet-img2img-layerprompt-checkbox { /*background-color: inherit; filter: contrast(1.5);*/ background-color: rgb(59 77 103 / var(--tw-bg-opacity));