added samples to img2img

fixed a bug with sampler selection (oops)

1 files changed, 10 insertions, 8 deletions

diff --git a/webui.py b/webui.py
index 03675032..4ff9d005 100644
--- a/webui.py
+++ b/webui.py
@@ -64,7 +64,7 @@ css_hide_progressbar = """
 

 SamplerData = namedtuple('SamplerData', ['name', 'constructor'])

 samplers = [

-    *[SamplerData(x[0], lambda model: KDiffusionSampler(model, x[1])) for x in [

+    *[SamplerData(x[0], lambda m, funcname=x[1]: KDiffusionSampler(m, funcname)) for x in [

         ('LMS', 'sample_lms'),

         ('Heun', 'sample_heun'),

         ('Euler', 'sample_euler'),

@@ -72,9 +72,10 @@ samplers = [
         ('DPM 2', 'sample_dpm_2'),

         ('DPM 2 Ancestral', 'sample_dpm_2_ancestral'),

     ] if hasattr(k_diffusion.sampling, x[1])],

-    SamplerData('DDIM', lambda model: DDIMSampler(model)),

-    SamplerData('PLMS', lambda model: PLMSSampler(model)),

+    SamplerData('DDIM', lambda m: DDIMSampler(model)),

+    SamplerData('PLMS', lambda m: PLMSSampler(model)),

 ]

+samplers_for_img2img = [x for x in samplers if x.name != 'DDIM' and x.name != 'PLMS']

 

 RealesrganModelInfo = namedtuple("RealesrganModelInfo", ["name", "location", "model", "netscale"])

 

@@ -197,14 +198,14 @@ class KDiffusionSampler:
         self.model = m

         self.model_wrap = k_diffusion.external.CompVisDenoiser(m)

         self.funcname = funcname

+        self.func = getattr(k_diffusion.sampling, self.funcname)

 

     def sample(self, S, conditioning, batch_size, shape, verbose, unconditional_guidance_scale, unconditional_conditioning, eta, x_T):

         sigmas = self.model_wrap.get_sigmas(S)

         x = x_T * sigmas[0]

         model_wrap_cfg = CFGDenoiser(self.model_wrap)

 

-        fun = getattr(k_diffusion.sampling, self.funcname)

-        samples_ddim = fun(model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}, disable=False)

+        samples_ddim = self.func(model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}, disable=False)

 

         return samples_ddim, None

 

@@ -810,10 +811,10 @@ txt2img_interface = gr.Interface(
 )

 

 

-def img2img(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, prompt_matrix, loopback: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int):

+def img2img(prompt: str, init_img, ddim_steps: int, sampler_index: int, use_GFPGAN: bool, prompt_matrix, loopback: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int):

     outpath = opts.outdir or "outputs/img2img-samples"

 

-    sampler = KDiffusionSampler(model, 'sample_lms')

+    sampler = samplers_for_img2img[sampler_index].constructor(model)

 

     assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'

 

@@ -842,7 +843,7 @@ def img2img(prompt: str, init_img, ddim_steps: int, use_GFPGAN: bool, prompt_mat
         xi = x0 + noise

         sigma_sched = sigmas[ddim_steps - t_enc - 1:]

         model_wrap_cfg = CFGDenoiser(sampler.model_wrap)

-        samples_ddim = k_diffusion.sampling.sample_lms(model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': cfg_scale}, disable=False)

+        samples_ddim = sampler.func(model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': cfg_scale}, disable=False)

         return samples_ddim

 

     if loopback:

@@ -919,6 +920,7 @@ img2img_interface = gr.Interface(
         gr.Textbox(placeholder="A fantasy landscape, trending on artstation.", lines=1),

         gr.Image(value=sample_img2img, source="upload", interactive=True, type="pil"),

         gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=50),

+        gr.Radio(label='Sampling method', choices=[x.name for x in samplers_for_img2img], value=samplers_for_img2img[0].name, type="index"),

         gr.Checkbox(label='Fix faces using GFPGAN', value=False, visible=GFPGAN is not None),

         gr.Checkbox(label='Create prompt matrix (separate multiple prompts using |, and get all combinations of them)', value=False),

         gr.Checkbox(label='Loopback (use images from previous batch when creating next batch)', value=False),