mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
41 lines
2.0 KiB
Python
41 lines
2.0 KiB
Python
import gradio as gr
|
|
import diffusers
|
|
from modules import scripts_manager, processing, shared, sd_models, devices
|
|
|
|
|
|
class Script(scripts_manager.Script):
|
|
def title(self):
|
|
return 'Kohya HiRes Fix'
|
|
|
|
def show(self, is_img2img):
|
|
return not is_img2img
|
|
|
|
# return signature is array of gradio components
|
|
def ui(self, _is_img2img):
|
|
with gr.Row():
|
|
gr.HTML('<a href="https://github.com/huggingface/diffusers/pull/7633">  Kohya HiRes Fix</a><br>')
|
|
with gr.Row():
|
|
enabled = gr.Checkbox(label="Enabled", value=True)
|
|
with gr.Row():
|
|
scale_factor = gr.Slider(value=0.5, minimum=0, maximum=1, step=0.05, label="Scale factor")
|
|
timestep = gr.Number(value=600, minimum=0, maximum=1000, label="Timestep")
|
|
block_num = gr.Number(value=1, minimum=0, maximum=10, label="Block")
|
|
return [enabled, scale_factor, timestep, block_num]
|
|
|
|
def run(self, p: processing.StableDiffusionProcessing, enabled, scale_factor, timestep, block_num): # pylint: disable=arguments-differ
|
|
if not enabled:
|
|
return None
|
|
if shared.sd_model_type != 'sd':
|
|
shared.log.warning(f'Kohya Hires Fix: pipeline={shared.sd_model_type} required=sd')
|
|
return None
|
|
old_pipe = shared.sd_model
|
|
high_res_fix = [{'timestep': timestep, 'scale_factor': scale_factor, 'block_num': block_num}]
|
|
shared.sd_model = diffusers.StableDiffusionPipeline.from_pipe(shared.sd_model, **{ 'custom_pipeline': 'kohya_hires_fix', 'high_res_fix': high_res_fix })
|
|
sd_models.copy_diffuser_options(shared.sd_model, old_pipe)
|
|
sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
|
|
sd_models.set_diffuser_options(shared.sd_model, vae=None, op='model')
|
|
shared.log.debug(f'Kohya Hires Fix: pipeline={shared.sd_model.__class__.__name__} args={high_res_fix}')
|
|
processed = processing.process_images(p)
|
|
shared.sd_model = old_pipe
|
|
return processed
|