分开两个tag,增加examples

app.py

@@ -419,7 +419,10 @@ def merge_images(images: np.ndarray):
 
 def get_models(device):
     def modelLoad(model, model_path, data_parallel=False):
-        model.load_state_dict(torch.load(model_path, map_location='cpu'), strict=True)
+        if str(device) == "cuda":
+            model.load_state_dict(torch.load(model_path), strict=True)
+        else:
+            model.load_state_dict(torch.load(model_path, map_location='cpu'), strict=True)
 
         if data_parallel:
             model = torch.nn.DataParallel(model)
@@ -461,10 +464,11 @@ def get_models(device):
 
 def init_webui(unet, vae, normal_t):
     # 定义回调函数
-    def process_image(input_image_value, noise_step, step_value, batch_size, sampler_name, img_size,
+    def process_image(input_image_value, noise_step, step_value, batch_size, sampler_name, img_size, random_seed,
                       progress=gr.Progress()):
         progress(0, desc="开始...")
 
+        setup_seed(int(random_seed))
         noise_step = float(noise_step)
         step_value = int(step_value)
         batch_size = int(batch_size)
@@ -481,7 +485,8 @@ def init_webui(unet, vae, normal_t):
         if input_image_value is None:
             looper = sampler.sample_loop(unet, vae.middle_c, batch_size, step_value, shape=img_size, eta=1.)
         else:
-            input_image_value = Image.fromarray(input_image_value).resize(img_size, Image.ANTIALIAS)
+            input_image_value = Image.fromarray(input_image_value).resize((img_size[0] * 8, img_size[1] * 8),
+                                                                          Image.ANTIALIAS)
             input_image_value = np.array(input_image_value, dtype=np.float32) / 255.
             input_image_value = np.transpose(input_image_value, (2, 0, 1))
             input_image_value = torch.Tensor([input_image_value]).to(device)
@@ -504,34 +509,93 @@ def init_webui(unet, vae, normal_t):
 
         return output
 
-    with gr.Blocks(title="图片处理") as iface:
-        with gr.Column():
-            with gr.Row():
-                with gr.Column():
-                    # 创建输入组件
-                    input_image = gr.Image(label="输入图片")
-                    # 加噪程度
-                    noise_step = gr.Slider(minimum=0.05, maximum=1, value=0.6, label="加噪程度", step=0.01)
-                with gr.Column():
+    def process_image_u(step_value, batch_size, sampler_name, img_size, random_seed,
+                        progress=gr.Progress()):
+        return process_image(None, 0, step_value, batch_size, sampler_name, img_size, random_seed,
+                             progress)
+
+    with gr.Blocks() as iface:
+        gr.Markdown(
+            "This is a diffusion model for generating second-dimensional avatars, which can be used for unconditional generation or image-to-image generation.")
+
+        with gr.Tab(label="unconditional generation"):
+            with gr.Column():
+                with gr.Row():
                     # 选择sampler
-                    sampler_name = gr.Dropdown(["DDIM"], label="sampler", value="DDIM")  # , "euler a", "dpmpp 2m"
+                    sampler_name_u = gr.Dropdown(["DDIM"], label="sampler", value="DDIM")  # , "euler a", "dpmpp 2m"
                     # 创建滑动条组件
-                    step = gr.Slider(minimum=1, maximum=1000, value=40, label="步长", step=1)
-                    batch_size = gr.Slider(minimum=1, maximum=4, label="batch size", step=1)
-                    img_size = gr.Slider(minimum=256, maximum=512, value=256, label="img size", step=64)
+                    step_u = gr.Slider(minimum=1, maximum=1000, value=40, label="steps", step=1)
+                    batch_size_u = gr.Slider(minimum=1, maximum=4, label="batch size", step=1)
+                    img_size_u = gr.Slider(minimum=256, maximum=512, value=256, label="img size", step=64)
+                    ramdom_seed_u = gr.Number(value=-1, label="ramdom seed(-1 as random number)")
                     # 创建开始按钮组件
-                    start_button = gr.Button(label="开始")
+                    start_button_u = gr.Button(label="Run")
+            # 创建输出组件
+            output_images_u = gr.Gallery(show_label=False, height=400, columns=5)
+            gr.Examples(
+                examples=[["DDIM", 40, 2, 256, 255395]],
+                inputs=[sampler_name_u, step_u, batch_size_u, img_size_u, ramdom_seed_u],
+                outputs=output_images_u,
+                fn=process_image_u,
+                # cache_examples=True,
+            )
+        with gr.Tab(label="image to image"):
+            with gr.Column():
+                with gr.Row():
+                    with gr.Column():
+                        # 创建输入组件
+                        input_image = gr.Image(label="image to image")
+                        # 加噪程度
+                        noise_step = gr.Slider(minimum=0.05, maximum=1, value=0.6, label="加噪程度", step=0.01)
+                    with gr.Column():
+                        # 选择sampler
+                        sampler_name = gr.Dropdown(["DDIM"], label="sampler", value="DDIM")  # , "euler a", "dpmpp 2m"
+                        # 创建滑动条组件
+                        step = gr.Slider(minimum=1, maximum=1000, value=40, label="steps", step=1)
+                        batch_size = gr.Slider(minimum=1, maximum=4, label="batch size", step=1)
+                        img_size = gr.Slider(minimum=256, maximum=512, value=256, label="img size", step=64)
+                        ramdom_seed = gr.Number(value=-1, label="ramdom seed(-1 as random number)")
+                        # 创建开始按钮组件
+                        start_button = gr.Button(label="Run")
+
             # 创建输出组件
             output_images = gr.Gallery(show_label=False, height=400, columns=5)
+            gr.Examples(
+                examples=[["./example.jpg", 0.4, "DDIM", 60, 4, 320, 231324]],  # 224477,378754
+                inputs=[input_image, noise_step, sampler_name, step, batch_size, img_size, ramdom_seed],
+                outputs=output_images,
+                fn=process_image,
+                # cache_examples=True,
+            )
 
-        start_button.click(process_image, [input_image, noise_step, step, batch_size, sampler_name, img_size],
+        start_button.click(process_image,
+                           [input_image, noise_step, step, batch_size, sampler_name, img_size, ramdom_seed],
                            [output_images])
+        start_button_u.click(process_image_u, [step_u, batch_size_u, sampler_name_u, img_size_u, ramdom_seed_u],
+                             [output_images_u])
 
     return iface
 
 
+def setup_seed(seed=0):
+    import random
+    if seed == -1:
+        seed = random.randint(0, 1000000)
+    print(seed)
+    torch.manual_seed(seed)  # 为CPU设置随机种子
+    np.random.seed(seed)  # Numpy module.
+    random.seed(seed)  # Python random module.
+    if torch.cuda.is_available():
+        # torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+        torch.cuda.manual_seed(seed)  # 为当前GPU设置随机种子
+        torch.cuda.manual_seed_all(seed)  # 为所有GPU设置随机种子
+        # os.environ['PYTHONHASHSEED'] = str(seed)
+
+
 if __name__ == '__main__':
     device = torch.device('cpu')
+    # device = torch.device('cuda')
     unet, vae, normal_t = get_models(device)
 
 

requirements.txt

@@ -5,4 +5,5 @@ numpy==1.23.4
 torch
 tqdm
 diffusers==0.20.1
-accelerate
+accelerate
+Pillow==9.2.0