upload upscale model

diffusion_webui/__init__.py

@@ -1 +1 @@
-__version__ = "1.6.2"
+__version__ = "1.8.0"

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_canny.py

@@ -3,9 +3,11 @@ import gradio as gr
 import numpy as np
 import torch
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_canny_model_list,
     stable_inpiant_model_list,
@@ -27,11 +29,13 @@ class StableDiffusionControlNetInpaintCannyGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -39,7 +43,7 @@ class StableDiffusionControlNetInpaintCannyGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-    
+
     def load_image(self, image_path):
         image = np.array(image_path)
         image = Image.fromarray(image)
@@ -76,10 +80,10 @@ class StableDiffusionControlNetInpaintCannyGenerator:
 
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         control_image = self.controlnet_canny_inpaint(image_path=image_path)
         pipe = self.load_model(
             stable_model_path=stable_model_path,

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_depth.py

@@ -2,10 +2,12 @@ import gradio as gr
 import numpy as np
 import torch
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 from transformers import pipeline
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_depth_model_list,
     stable_inpiant_model_list,
@@ -27,11 +29,13 @@ class StableDiffusionControlInpaintNetDepthGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -72,10 +76,10 @@ class StableDiffusionControlInpaintNetDepthGenerator:
     ):
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         control_image = self.controlnet_inpaint_depth(image_path=image_path)
 
         pipe = self.load_model(
@@ -92,7 +96,6 @@ class StableDiffusionControlInpaintNetDepthGenerator:
 
         output = pipe(
             prompt=prompt,
-            
             image=normal_image,
             mask_image=mask_image,
             control_image=control_image,

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_hed.py

@@ -3,8 +3,11 @@ import numpy as np
 import torch
 from controlnet_aux import HEDdetector
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
+from PIL import Image
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_hed_model_list,
     stable_inpiant_model_list,
@@ -13,7 +16,6 @@ from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
 )
-from PIL import Image
 
 # https://github.com/mikonvergence/ControlNetInpaint
 
@@ -27,11 +29,13 @@ class StableDiffusionControlNetInpaintHedGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -45,7 +49,6 @@ class StableDiffusionControlNetInpaintHedGenerator:
         image = Image.fromarray(image)
         return image
 
-
     def controlnet_inpaint_hed(self, image_path: str):
         hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
         image = image_path["image"].convert("RGB").resize((512, 512))
@@ -70,10 +73,10 @@ class StableDiffusionControlNetInpaintHedGenerator:
     ):
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         control_image = self.controlnet_inpaint_hed(image_path=image_path)
 
         pipe = self.load_model(

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_mlsd.py

@@ -3,8 +3,11 @@ import numpy as np
 import torch
 from controlnet_aux import MLSDdetector
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
+from PIL import Image
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_mlsd_model_list,
     stable_inpiant_model_list,
@@ -13,7 +16,6 @@ from diffusion_webui.utils.scheduler_list import (
     SCHEDULER_LIST,
     get_scheduler_list,
 )
-from PIL import Image
 
 # https://github.com/mikonvergence/ControlNetInpaint
 
@@ -27,11 +29,13 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
             controlnet = ControlNetModel.from_pretrained(
                 controlnet_model_path, torch_dtype=torch.float16
             )
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -39,13 +43,12 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-    
+
     def load_image(self, image_path):
         image = np.array(image_path)
         image = Image.fromarray(image)
         return image
 
-
     def controlnet_inpaint_mlsd(self, image_path: str):
         mlsd = MLSDdetector.from_pretrained("lllyasviel/ControlNet")
         image = image_path["image"].convert("RGB").resize((512, 512))
@@ -71,10 +74,10 @@ class StableDiffusionControlNetInpaintMlsdGenerator:
 
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         control_image = self.controlnet_inpaint_mlsd(image_path=image_path)
 
         pipe = self.load_model(

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_pose.py

@@ -3,9 +3,11 @@ import numpy as np
 import torch
 from controlnet_aux import OpenposeDetector
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_pose_model_list,
     stable_inpiant_model_list,
@@ -28,11 +30,13 @@ class StableDiffusionControlNetInpaintPoseGenerator:
                 controlnet_model_path, torch_dtype=torch.float16
             )
 
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -40,12 +44,12 @@ class StableDiffusionControlNetInpaintPoseGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-    
+
     def load_image(self, image_path):
         image = np.array(image_path)
         image = Image.fromarray(image)
         return image
-    
+
     def controlnet_pose_inpaint(self, image_path: str):
         openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
@@ -71,10 +75,10 @@ class StableDiffusionControlNetInpaintPoseGenerator:
     ):
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         controlnet_image = self.controlnet_pose_inpaint(image_path=image_path)
 
         pipe = self.load_model(

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_scribble.py

@@ -3,9 +3,11 @@ import numpy as np
 import torch
 from controlnet_aux import HEDdetector
 from diffusers import ControlNetModel
-from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline
 from PIL import Image
 
+from diffusion_webui.diffusion_models.controlnet.controlnet_inpaint.pipeline_stable_diffusion_controlnet_inpaint import (
+    StableDiffusionControlNetInpaintPipeline,
+)
 from diffusion_webui.utils.model_list import (
     controlnet_scribble_model_list,
     stable_inpiant_model_list,
@@ -17,6 +19,7 @@ from diffusion_webui.utils.scheduler_list import (
 
 # https://github.com/mikonvergence/ControlNetInpaint
 
+
 class StableDiffusionControlNetInpaintScribbleGenerator:
     def __init__(self):
         self.pipe = None
@@ -27,11 +30,13 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
                 controlnet_model_path, torch_dtype=torch.float16
             )
 
-            self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-                pretrained_model_name_or_path=stable_model_path,
-                controlnet=controlnet,
-                safety_checker=None,
-                torch_dtype=torch.float16,
+            self.pipe = (
+                StableDiffusionControlNetInpaintPipeline.from_pretrained(
+                    pretrained_model_name_or_path=stable_model_path,
+                    controlnet=controlnet,
+                    safety_checker=None,
+                    torch_dtype=torch.float16,
+                )
             )
 
         self.pipe = get_scheduler_list(pipe=self.pipe, scheduler=scheduler)
@@ -39,7 +44,7 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-    
+
     def load_image(self, image_path):
         image = np.array(image_path)
         image = Image.fromarray(image)
@@ -70,11 +75,13 @@ class StableDiffusionControlNetInpaintScribbleGenerator:
     ):
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
-        controlnet_image = self.controlnet_inpaint_scribble(image_path=image_path)
+
+        controlnet_image = self.controlnet_inpaint_scribble(
+            image_path=image_path
+        )
 
         pipe = self.load_model(
             stable_model_path=stable_model_path,

diffusion_webui/diffusion_models/controlnet/controlnet_inpaint/controlnet_inpaint_seg.py

@@ -200,11 +200,12 @@ class StableDiffusionControlNetInpaintSegGenerator:
         self.pipe.enable_xformers_memory_efficient_attention()
 
         return self.pipe
-    
+
     def load_image(self, image_path):
         image = np.array(image_path)
         image = Image.fromarray(image)
         return image
+
     def controlnet_seg_inpaint(self, image_path: str):
         image_processor = AutoImageProcessor.from_pretrained(
             "openmmlab/upernet-convnext-small"
@@ -252,10 +253,10 @@ class StableDiffusionControlNetInpaintSegGenerator:
 
         normal_image = image_path["image"].convert("RGB").resize((512, 512))
         mask_image = image_path["mask"].convert("RGB").resize((512, 512))
-        
+
         normal_image = self.load_image(image_path=normal_image)
         mask_image = self.load_image(image_path=mask_image)
-        
+
         controlnet_image = self.controlnet_seg_inpaint(image_path=image_path)
 
         pipe = self.load_model(

diffusion_webui/helpers.py

@@ -49,3 +49,6 @@ from diffusion_webui.diffusion_models.stable_diffusion.inpaint_app import (
 from diffusion_webui.diffusion_models.stable_diffusion.text2img_app import (
     StableDiffusionText2ImageGenerator,
 )
+from diffusion_webui.upscaler_models.codeformer_upscaler import (
+    CodeformerUpscalerGenerator,
+)

diffusion_webui/upscaler_models/codeformer_upscaler.py

@@ -0,0 +1,84 @@
+import gradio as gr
+from codeformer.app import inference_app
+
+
+class CodeformerUpscalerGenerator:
+    def __init__(self):
+        self.pipe = None
+
+    def generate_image(
+        self,
+        image_path: str,
+        background_enhance: bool,
+        face_upsample: bool,
+        upscale: int,
+        codeformer_fidelity: int,
+    ):
+        if self.pipe is None:
+            self.pipe = inference_app(
+                image=image_path,
+                background_enhance=background_enhance,
+                face_upsample=face_upsample,
+                upscale=upscale,
+                codeformer_fidelity=codeformer_fidelity,
+            )
+
+        return [self.pipe]
+
+    def app():
+        with gr.Blocks():
+            with gr.Row():
+                with gr.Column():
+                    codeformer_upscale_image_file = gr.Image(
+                        type="filepath", label="Image"
+                    ).style(height=260)
+
+                    with gr.Row():
+                        with gr.Column():
+                            codeformer_face_upsample = gr.Checkbox(
+                                label="Face Upsample",
+                                value=True,
+                            )
+                            codeformer_upscale = gr.Slider(
+                                label="Upscale",
+                                minimum=1,
+                                maximum=4,
+                                step=1,
+                                value=2,
+                            )
+                        with gr.Row():
+                            with gr.Column():
+                                codeformer_background_enhance = gr.Checkbox(
+                                    label="Background Enhance",
+                                    value=True,
+                                )
+                                codeformer_upscale_fidelity = gr.Slider(
+                                    label="Codeformer Fidelity",
+                                    minimum=0.1,
+                                    maximum=1.0,
+                                    step=0.1,
+                                    value=0.5,
+                                )
+
+                    codeformer_upscale_predict_button = gr.Button(
+                        value="Generator"
+                    )
+
+                with gr.Column():
+                    output_image = gr.Gallery(
+                        label="Generated images",
+                        show_label=False,
+                        elem_id="gallery",
+                    ).style(grid=(1, 2))
+
+        codeformer_upscale_predict_button.click(
+            fn=CodeformerUpscalerGenerator().generate_image,
+            inputs=[
+                codeformer_upscale_image_file,
+                codeformer_background_enhance,
+                codeformer_face_upsample,
+                codeformer_upscale,
+                codeformer_upscale_fidelity,
+            ],
+            outputs=[output_image],
+        )

diffusion_webui/utils/model_list.py

@@ -29,8 +29,8 @@ controlnet_scribble_model_list = [
     "thibaud/controlnet-sd21-scribble-diffusers",
 ]
 stable_inpiant_model_list = [
-    "runwayml/stable-diffusion-inpainting",
     "stabilityai/stable-diffusion-2-inpainting",
+    "runwayml/stable-diffusion-inpainting",
 ]
 
 controlnet_mlsd_model_list = [