From 8bdb8f688972aa48b98fa8e4eec6a6803a01d866 Mon Sep 17 00:00:00 2001
From: lvmin <lyuminzhang@outlook.com>
Date: Thu, 10 Aug 2023 21:00:55 -0700
Subject: [PATCH] i

---
 modules/default_pipeline.py | 12 ++++++------
 webui.py                    | 31 ++++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/modules/default_pipeline.py b/modules/default_pipeline.py
index 56dc023f..0bd951a9 100644
--- a/modules/default_pipeline.py
+++ b/modules/default_pipeline.py
@@ -18,14 +18,14 @@ xl_refiner = core.load_model(xl_refiner_filename)
 
 
 @torch.no_grad()
-def process(positive_prompt, negative_prompt, width=1280, height=960, batch_size=1):
+def process(positive_prompt, negative_prompt, steps, switch, width, height, image_seed):
     positive_conditions = core.encode_prompt_condition(clip=xl_base.clip, prompt=positive_prompt)
     negative_conditions = core.encode_prompt_condition(clip=xl_base.clip, prompt=negative_prompt)
 
     positive_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=positive_prompt)
     negative_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=negative_prompt)
 
-    empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=batch_size)
+    empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
 
     sampled_latent = core.ksampler_with_refiner(
         model=xl_base.unet,
@@ -34,10 +34,10 @@ def process(positive_prompt, negative_prompt, width=1280, height=960, batch_size
         refiner=xl_refiner.unet,
         refiner_positive=positive_conditions_refiner,
         refiner_negative=negative_conditions_refiner,
-        refiner_switch_step=20,
+        refiner_switch_step=switch,
         latent=empty_latent,
-        steps=30, start_step=0, last_step=30, disable_noise=False, force_full_denoise=True,
-        seed=123456
+        steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
+        seed=image_seed
     )
 
     decoded_latent = core.decode_vae(vae=xl_refiner.vae, latent_image=sampled_latent)
@@ -46,4 +46,4 @@ def process(positive_prompt, negative_prompt, width=1280, height=960, batch_size
 
     close_all_preview()
 
-    return images * 2
+    return images
diff --git a/webui.py b/webui.py
index 6b6e574e..8d01b47f 100644
--- a/webui.py
+++ b/webui.py
@@ -1,18 +1,35 @@
 import gradio as gr
+import random
 
 from modules.sdxl_styles import apply_style, style_keys, aspect_ratios
-# from modules.default_pipeline import process
+from modules.default_pipeline import process
 
 
-def generate_clicked(positive_prompt):
+def generate_clicked(prompt, negative_prompt, style_selction, performance_selction,
+                     aspect_ratios_selction, image_number, image_seed):
 
-    p, n = apply_style('cinematic-default', positive_prompt, '')
+    p_txt, n_txt = apply_style(style_selction, prompt, negative_prompt)
 
-    print(p)
-    print(n)
+    if performance_selction == 'Speed':
+        steps = 30
+        switch = 20
+    else:
+        steps = 60
+        switch = 40
 
-    return process(positive_prompt=p,
-                   negative_prompt=n)
+    width, height = aspect_ratios[aspect_ratios_selction]
+
+    results = []
+    seed = image_seed
+    if not isinstance(seed, int) or seed < 0 or seed > 65535:
+        seed = random.randint(1, 65535)
+
+    for i in range(image_number):
+        imgs = process(p_txt, n_txt, steps, switch, width, height, seed)
+        seed += 1
+        results += imgs
+
+    return results
 
 
 block = gr.Blocks()