restructure for gpu

app.py

@@ -1,3 +1,6 @@
+import subprocess
+import pip
+
 import io
 import base64
 import os
@@ -86,15 +89,40 @@ USE_GLID = False
 # except:
 #     USE_GLID = False
 
-try:
-    import onnxruntime
-    onnx_available = True
-    onnx_providers = ["CUDAExecutionProvider", "DmlExecutionProvider", "OpenVINOExecutionProvider", 'CPUExecutionProvider']
-    available_providers = onnxruntime.get_available_providers()
-    onnx_providers = [item for item in onnx_providers if item in available_providers]
-except:
-    onnx_available = False
-    onnx_providers = []
+# ******** ORIGINAL ***********
+# try:
+#     import onnxruntime
+#     onnx_available = True
+#     onnx_providers = ["CUDAExecutionProvider", "DmlExecutionProvider", "OpenVINOExecutionProvider", 'CPUExecutionProvider']
+#     available_providers = onnxruntime.get_available_providers()
+#     onnx_providers = [item for item in onnx_providers if item in available_providers]
+# except:
+#     onnx_available = False
+#     onnx_providers = []
+
+
+# try:
+#     cuda_available = torch.cuda.is_available()
+# except:
+#     cuda_available = False
+# finally:
+#     if sys.platform == "darwin":
+#         device = "mps" if torch.backends.mps.is_available() else "cpu"
+#     elif cuda_available:
+#         device = "cuda"
+#     else:
+#         device = "cpu"
+
+# if device != "cuda":
+#     import contextlib
+
+#     autocast = contextlib.nullcontext
+
+# with open("config.yaml", "r") as yaml_in:
+#     yaml_object = yaml.safe_load(yaml_in)
+#     config_json = json.dumps(yaml_object)
+    
+# ******** ^ ORIGINAL ^ ***********
 
 try:
     cuda_available = torch.cuda.is_available()
@@ -108,17 +136,17 @@ finally:
     else:
         device = "cpu"
 
-if device != "cuda":
-    import contextlib
+import contextlib
 
-    autocast = contextlib.nullcontext
+autocast = contextlib.nullcontext
 
 with open("config.yaml", "r") as yaml_in:
     yaml_object = yaml.safe_load(yaml_in)
     config_json = json.dumps(yaml_object)
-    
 
 
+#  new ^ 
+
 def load_html():
     body, canvaspy = "", ""
     with open("index.html", encoding="utf8") as f:
@@ -315,12 +343,13 @@ class StableDiffusionInpaint:
         vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
         # if device == "cuda" and not args.fp32:
         #     vae.to(torch.float16)
+        vae.to(torch.float16)
         if original_checkpoint:
             print(f"Converting & Loading {model_path}")
             from convert_checkpoint import convert_checkpoint
 
             pipe = convert_checkpoint(model_path, inpainting=True)
-            if device == "cuda" and not args.fp32:
+            if device == "cuda":
                 pipe.to(torch.float16)
             inpaint = StableDiffusionInpaintPipeline(
                 vae=vae,
@@ -333,7 +362,7 @@ class StableDiffusionInpaint:
             )
         else:
             print(f"Loading {model_name}")
-            if device == "cuda" and not args.fp32:
+            if device == "cuda":
                 inpaint = StableDiffusionInpaintPipeline.from_pretrained(
                     model_name,
                     revision="fp16",
@@ -476,7 +505,6 @@ class StableDiffusionInpaint:
                 )["images"]
         return images
 
-
 class StableDiffusion:
     def __init__(
         self,
@@ -488,134 +516,74 @@ class StableDiffusion:
     ):
         self.token = token
         original_checkpoint = False
-        if device=="cpu" and onnx_available:
-            from diffusers import OnnxStableDiffusionPipeline, OnnxStableDiffusionInpaintPipelineLegacy, OnnxStableDiffusionImg2ImgPipeline
-            text2img = OnnxStableDiffusionPipeline.from_pretrained(
-                model_name,
-                revision="onnx",
-                provider=onnx_providers[0] if onnx_providers else None
-                )
-            inpaint = OnnxStableDiffusionInpaintPipelineLegacy(
-                    vae_encoder=text2img.vae_encoder,
-                    vae_decoder=text2img.vae_decoder,
-                    text_encoder=text2img.text_encoder,
-                    tokenizer=text2img.tokenizer,
-                    unet=text2img.unet,
-                    scheduler=text2img.scheduler,
-                    safety_checker=text2img.safety_checker,
-                    feature_extractor=text2img.feature_extractor,
-                )
-            img2img = OnnxStableDiffusionImg2ImgPipeline(
-                vae_encoder=text2img.vae_encoder,
-                vae_decoder=text2img.vae_decoder,
-                text_encoder=text2img.text_encoder,
-                tokenizer=text2img.tokenizer,
-                unet=text2img.unet,
-                scheduler=text2img.scheduler,
-                safety_checker=text2img.safety_checker,
-                feature_extractor=text2img.feature_extractor,
-            )
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+        vae.to(torch.float16)
+        if model_path and os.path.exists(model_path):
+            if model_path.endswith(".ckpt"):
+                original_checkpoint = True
+            elif model_path.endswith(".json"):
+                model_name = os.path.dirname(model_path)
+            else:
+                model_name = model_path
+        if original_checkpoint:
+            print(f"Converting & Loading {model_path}")
+            from convert_checkpoint import convert_checkpoint
+
+            text2img = convert_checkpoint(model_path)
+            if device == "cuda" and not args.fp32:
+                text2img.to(torch.float16)
         else:
-            if model_path and os.path.exists(model_path):
-                if model_path.endswith(".ckpt"):
-                    original_checkpoint = True
-                elif model_path.endswith(".json"):
-                    model_name = os.path.dirname(model_path)
-                else:
-                    model_name = model_path
-            vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+            print(f"Loading {model_name}")
             if device == "cuda" and not args.fp32:
-                vae.to(torch.float16)
-            if original_checkpoint:
-                print(f"Converting & Loading {model_path}")
-                from convert_checkpoint import convert_checkpoint
-
-                pipe = convert_checkpoint(model_path)
-                if device == "cuda" and not args.fp32:
-                    pipe.to(torch.float16)
-                text2img = StableDiffusionPipeline(
-                    vae=vae,
-                    text_encoder=pipe.text_encoder,
-                    tokenizer=pipe.tokenizer,
-                    unet=pipe.unet,
-                    scheduler=pipe.scheduler,
-                    safety_checker=pipe.safety_checker,
-                    feature_extractor=pipe.feature_extractor,
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5",
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                    vae=vae
                 )
             else:
-                print(f"Loading {model_name}")
-                if device == "cuda" and not args.fp32:
-                    text2img = StableDiffusionPipeline.from_pretrained(
-                        model_name,
-                        revision="fp16",
-                        torch_dtype=torch.float16,
-                        use_auth_token=token,
-                        vae=vae,
-                    )
-                else:
-                    text2img = StableDiffusionPipeline.from_pretrained(
-                        model_name, use_auth_token=token, vae=vae
-                    )
-            if inpainting_model:
-                # can reduce vRAM by reusing models except unet
-                text2img_unet = text2img.unet
-                del text2img.vae
-                del text2img.text_encoder
-                del text2img.tokenizer
-                del text2img.scheduler
-                del text2img.safety_checker
-                del text2img.feature_extractor
-                import gc
-
-                gc.collect()
-                if device == "cuda" and not args.fp32:
-                    inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-                        "runwayml/stable-diffusion-inpainting",
-                        revision="fp16",
-                        torch_dtype=torch.float16,
-                        use_auth_token=token,
-                        vae=vae,
-                    ).to(device)
-                else:
-                    inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-                        "runwayml/stable-diffusion-inpainting",
-                        use_auth_token=token,
-                        vae=vae,
-                    ).to(device)
-                text2img_unet.to(device)
-                text2img = StableDiffusionPipeline(
-                    vae=inpaint.vae,
-                    text_encoder=inpaint.text_encoder,
-                    tokenizer=inpaint.tokenizer,
-                    unet=text2img_unet,
-                    scheduler=inpaint.scheduler,
-                    safety_checker=inpaint.safety_checker,
-                    feature_extractor=inpaint.feature_extractor,
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    model_name, use_auth_token=token,
                 )
+        if inpainting_model:
+            # can reduce vRAM by reusing models except unet
+            text2img_unet = text2img.unet
+            del text2img.vae
+            del text2img.text_encoder
+            del text2img.tokenizer
+            del text2img.scheduler
+            del text2img.safety_checker
+            del text2img.feature_extractor
+            import gc
+
+            gc.collect()
+            if device == "cuda":
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting",
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                    vae=vae
+                ).to(device)
             else:
-                inpaint = StableDiffusionInpaintPipelineLegacy(
-                    vae=text2img.vae,
-                    text_encoder=text2img.text_encoder,
-                    tokenizer=text2img.tokenizer,
-                    unet=text2img.unet,
-                    scheduler=text2img.scheduler,
-                    safety_checker=text2img.safety_checker,
-                    feature_extractor=text2img.feature_extractor,
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting", use_auth_token=token,
                 ).to(device)
-            text_encoder = text2img.text_encoder
-            tokenizer = text2img.tokenizer
-            if os.path.exists("./embeddings"):
-                for item in os.listdir("./embeddings"):
-                    if item.endswith(".bin"):
-                        load_learned_embed_in_clip(
-                            os.path.join("./embeddings", item),
-                            text2img.text_encoder,
-                            text2img.tokenizer,
-                        )
-            text2img.to(device)
-            if device == "mps":
-                _ = text2img("", num_inference_steps=1)
-            img2img = StableDiffusionImg2ImgPipeline(
+            text2img_unet.to(device)
+            del text2img
+            gc.collect()
+            text2img = StableDiffusionPipeline(
+                vae=inpaint.vae,
+                text_encoder=inpaint.text_encoder,
+                tokenizer=inpaint.tokenizer,
+                unet=text2img_unet,
+                scheduler=inpaint.scheduler,
+                safety_checker=inpaint.safety_checker,
+                feature_extractor=inpaint.feature_extractor,
+            )
+        else:
+            inpaint = StableDiffusionInpaintPipelineLegacy(
                 vae=text2img.vae,
                 text_encoder=text2img.text_encoder,
                 tokenizer=text2img.tokenizer,
@@ -624,6 +592,19 @@ class StableDiffusion:
                 safety_checker=text2img.safety_checker,
                 feature_extractor=text2img.feature_extractor,
             ).to(device)
+        text_encoder = text2img.text_encoder
+        tokenizer = text2img.tokenizer
+        if os.path.exists("./embeddings"):
+            for item in os.listdir("./embeddings"):
+                if item.endswith(".bin"):
+                    load_learned_embed_in_clip(
+                        os.path.join("./embeddings", item),
+                        text2img.text_encoder,
+                        text2img.tokenizer,
+                    )
+        text2img.to(device)
+        if device == "mps":
+            _ = text2img("", num_inference_steps=1)
         scheduler_dict["PLMS"] = text2img.scheduler
         scheduler_dict["DDIM"] = prepare_scheduler(
             DDIMScheduler(
@@ -639,44 +620,40 @@ class StableDiffusion:
                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
             )
         )
-        scheduler_dict["PNDM"] = prepare_scheduler(
-            PNDMScheduler(
-                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
-                skip_prk_steps=True
-            )
-        )
         scheduler_dict["DPM"] = prepare_scheduler(
             DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)
         )
         self.safety_checker = text2img.safety_checker
+        img2img = StableDiffusionImg2ImgPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
         save_token(token)
         try:
             total_memory = torch.cuda.get_device_properties(0).total_memory // (
                 1024 ** 3
             )
-            if total_memory <= 5 or args.lowvram:
+            if total_memory <= 5:
                 inpaint.enable_attention_slicing()
-                inpaint.enable_sequential_cpu_offload()
-                if inpainting_model:
-                    text2img.enable_attention_slicing()
-                    text2img.enable_sequential_cpu_offload()
         except:
             pass
         self.text2img = text2img
         self.inpaint = inpaint
         self.img2img = img2img
-        if True:
-            self.unified = inpaint
-        else:
-            self.unified = UnifiedPipeline(
-                vae=text2img.vae,
-                text_encoder=text2img.text_encoder,
-                tokenizer=text2img.tokenizer,
-                unet=text2img.unet,
-                scheduler=text2img.scheduler,
-                safety_checker=text2img.safety_checker,
-                feature_extractor=text2img.feature_extractor,
-            ).to(device)
+        self.unified = UnifiedPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
         self.inpainting_model = inpainting_model
 
     def run(
@@ -707,7 +684,7 @@ class StableDiffusion:
         selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
         for item in [text2img, inpaint, img2img, unified]:
             item.scheduler = selected_scheduler
-            if enable_safety or self.safety_checker is None:
+            if enable_safety:
                 item.safety_checker = self.safety_checker
             else:
                 item.safety_checker = lambda images, **kwargs: (images, False)
@@ -743,7 +720,7 @@ class StableDiffusion:
             if True:
                 images = img2img(
                     prompt=prompt,
-                    image=init_image.resize(
+                    init_image=init_image.resize(
                         (process_width, process_height), resample=SAMPLING_MODE
                     ),
                     strength=strength,
@@ -753,40 +730,33 @@ class StableDiffusion:
             if fill_mode == "g_diffuser" and not self.inpainting_model:
                 mask = 255 - mask
                 mask = mask[:, :, np.newaxis].repeat(3, axis=2)
-                img, mask = functbl[fill_mode](img, mask)
+                img, mask, out_mask = functbl[fill_mode](img, mask)
                 extra_kwargs["strength"] = 1.0
-                extra_kwargs["out_mask"] = Image.fromarray(mask)
+                extra_kwargs["out_mask"] = Image.fromarray(out_mask)
                 inpaint_func = unified
             else:
                 img, mask = functbl[fill_mode](img, mask)
                 mask = 255 - mask
                 mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
                 mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+                extra_kwargs["strength"] = strength
                 inpaint_func = inpaint
             init_image = Image.fromarray(img)
             mask_image = Image.fromarray(mask)
             # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
-            input_image = init_image.resize(
-                (process_width, process_height), resample=SAMPLING_MODE
-            )
-            if self.inpainting_model:
+            if True:
+                input_image = init_image.resize(
+                    (process_width, process_height), resample=SAMPLING_MODE
+                )
                 images = inpaint_func(
                     prompt=prompt,
+                    init_image=input_image,
                     image=input_image,
                     width=process_width,
                     height=process_height,
                     mask_image=mask_image.resize((process_width, process_height)),
                     **extra_kwargs,
                 )["images"]
-            else:
-                extra_kwargs["strength"] = strength
-                if True:
-                    images = inpaint_func(
-                        prompt=prompt,
-                        image=input_image,
-                        mask_image=mask_image.resize((process_width, process_height)),
-                        **extra_kwargs,
-                    )["images"]
         else:
             if True:
                 images = text2img(
@@ -798,6 +768,327 @@ class StableDiffusion:
         return images
 
 
+# class StableDiffusion:
+#     def __init__(
+#         self,
+#         token: str = "",
+#         model_name: str = "runwayml/stable-diffusion-v1-5",
+#         model_path: str = None,
+#         inpainting_model: bool = False,
+#         **kwargs,
+#     ):
+#         self.token = token
+#         original_checkpoint = False
+#         if device=="cpu" and onnx_available:
+#             from diffusers import OnnxStableDiffusionPipeline, OnnxStableDiffusionInpaintPipelineLegacy, OnnxStableDiffusionImg2ImgPipeline
+#             text2img = OnnxStableDiffusionPipeline.from_pretrained(
+#                 model_name,
+#                 revision="onnx",
+#                 provider=onnx_providers[0] if onnx_providers else None
+#                 )
+#             inpaint = OnnxStableDiffusionInpaintPipelineLegacy(
+#                     vae_encoder=text2img.vae_encoder,
+#                     vae_decoder=text2img.vae_decoder,
+#                     text_encoder=text2img.text_encoder,
+#                     tokenizer=text2img.tokenizer,
+#                     unet=text2img.unet,
+#                     scheduler=text2img.scheduler,
+#                     safety_checker=text2img.safety_checker,
+#                     feature_extractor=text2img.feature_extractor,
+#                 )
+#             img2img = OnnxStableDiffusionImg2ImgPipeline(
+#                 vae_encoder=text2img.vae_encoder,
+#                 vae_decoder=text2img.vae_decoder,
+#                 text_encoder=text2img.text_encoder,
+#                 tokenizer=text2img.tokenizer,
+#                 unet=text2img.unet,
+#                 scheduler=text2img.scheduler,
+#                 safety_checker=text2img.safety_checker,
+#                 feature_extractor=text2img.feature_extractor,
+#             )
+#         else:
+#             if model_path and os.path.exists(model_path):
+#                 if model_path.endswith(".ckpt"):
+#                     original_checkpoint = True
+#                 elif model_path.endswith(".json"):
+#                     model_name = os.path.dirname(model_path)
+#                 else:
+#                     model_name = model_path
+#             vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+#             if device == "cuda" and not args.fp32:
+#                 vae.to(torch.float16)
+#             if original_checkpoint:
+#                 print(f"Converting & Loading {model_path}")
+#                 from convert_checkpoint import convert_checkpoint
+
+#                 pipe = convert_checkpoint(model_path)
+#                 if device == "cuda" and not args.fp32:
+#                     pipe.to(torch.float16)
+#                 text2img = StableDiffusionPipeline(
+#                     vae=vae,
+#                     text_encoder=pipe.text_encoder,
+#                     tokenizer=pipe.tokenizer,
+#                     unet=pipe.unet,
+#                     scheduler=pipe.scheduler,
+#                     safety_checker=pipe.safety_checker,
+#                     feature_extractor=pipe.feature_extractor,
+#                 )
+#             else:
+#                 print(f"Loading {model_name}")
+#                 if device == "cuda" and not args.fp32:
+#                     text2img = StableDiffusionPipeline.from_pretrained(
+#                         model_name,
+#                         revision="fp16",
+#                         torch_dtype=torch.float16,
+#                         use_auth_token=token,
+#                         vae=vae,
+#                     )
+#                 else:
+#                     text2img = StableDiffusionPipeline.from_pretrained(
+#                         model_name, use_auth_token=token, vae=vae
+#                     )
+#             if inpainting_model:
+#                 # can reduce vRAM by reusing models except unet
+#                 text2img_unet = text2img.unet
+#                 del text2img.vae
+#                 del text2img.text_encoder
+#                 del text2img.tokenizer
+#                 del text2img.scheduler
+#                 del text2img.safety_checker
+#                 del text2img.feature_extractor
+#                 import gc
+
+#                 gc.collect()
+#                 if device == "cuda" and not args.fp32:
+#                     inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+#                         "runwayml/stable-diffusion-inpainting",
+#                         revision="fp16",
+#                         torch_dtype=torch.float16,
+#                         use_auth_token=token,
+#                         vae=vae,
+#                     ).to(device)
+#                 else:
+#                     inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+#                         "runwayml/stable-diffusion-inpainting",
+#                         use_auth_token=token,
+#                         vae=vae,
+#                     ).to(device)
+#                 text2img_unet.to(device)
+#                 text2img = StableDiffusionPipeline(
+#                     vae=inpaint.vae,
+#                     text_encoder=inpaint.text_encoder,
+#                     tokenizer=inpaint.tokenizer,
+#                     unet=text2img_unet,
+#                     scheduler=inpaint.scheduler,
+#                     safety_checker=inpaint.safety_checker,
+#                     feature_extractor=inpaint.feature_extractor,
+#                 )
+#             else:
+#                 inpaint = StableDiffusionInpaintPipelineLegacy(
+#                     vae=text2img.vae,
+#                     text_encoder=text2img.text_encoder,
+#                     tokenizer=text2img.tokenizer,
+#                     unet=text2img.unet,
+#                     scheduler=text2img.scheduler,
+#                     safety_checker=text2img.safety_checker,
+#                     feature_extractor=text2img.feature_extractor,
+#                 ).to(device)
+#             text_encoder = text2img.text_encoder
+#             tokenizer = text2img.tokenizer
+#             if os.path.exists("./embeddings"):
+#                 for item in os.listdir("./embeddings"):
+#                     if item.endswith(".bin"):
+#                         load_learned_embed_in_clip(
+#                             os.path.join("./embeddings", item),
+#                             text2img.text_encoder,
+#                             text2img.tokenizer,
+#                         )
+#             text2img.to(device)
+#             if device == "mps":
+#                 _ = text2img("", num_inference_steps=1)
+#             img2img = StableDiffusionImg2ImgPipeline(
+#                 vae=text2img.vae,
+#                 text_encoder=text2img.text_encoder,
+#                 tokenizer=text2img.tokenizer,
+#                 unet=text2img.unet,
+#                 scheduler=text2img.scheduler,
+#                 safety_checker=text2img.safety_checker,
+#                 feature_extractor=text2img.feature_extractor,
+#             ).to(device)
+#         scheduler_dict["PLMS"] = text2img.scheduler
+#         scheduler_dict["DDIM"] = prepare_scheduler(
+#             DDIMScheduler(
+#                 beta_start=0.00085,
+#                 beta_end=0.012,
+#                 beta_schedule="scaled_linear",
+#                 clip_sample=False,
+#                 set_alpha_to_one=False,
+#             )
+#         )
+#         scheduler_dict["K-LMS"] = prepare_scheduler(
+#             LMSDiscreteScheduler(
+#                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+#             )
+#         )
+#         scheduler_dict["PNDM"] = prepare_scheduler(
+#             PNDMScheduler(
+#                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
+#                 skip_prk_steps=True
+#             )
+#         )
+#         scheduler_dict["DPM"] = prepare_scheduler(
+#             DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)
+#         )
+#         self.safety_checker = text2img.safety_checker
+#         save_token(token)
+#         try:
+#             total_memory = torch.cuda.get_device_properties(0).total_memory // (
+#                 1024 ** 3
+#             )
+#             if total_memory <= 5 or args.lowvram:
+#                 inpaint.enable_attention_slicing()
+#                 inpaint.enable_sequential_cpu_offload()
+#                 if inpainting_model:
+#                     text2img.enable_attention_slicing()
+#                     text2img.enable_sequential_cpu_offload()
+#         except:
+#             pass
+#         self.text2img = text2img
+#         self.inpaint = inpaint
+#         self.img2img = img2img
+#         if True:
+#             self.unified = inpaint
+#         else:
+#             self.unified = UnifiedPipeline(
+#                 vae=text2img.vae,
+#                 text_encoder=text2img.text_encoder,
+#                 tokenizer=text2img.tokenizer,
+#                 unet=text2img.unet,
+#                 scheduler=text2img.scheduler,
+#                 safety_checker=text2img.safety_checker,
+#                 feature_extractor=text2img.feature_extractor,
+#             ).to(device)
+#         self.inpainting_model = inpainting_model
+
+#     def run(
+#         self,
+#         image_pil,
+#         prompt="",
+#         negative_prompt="",
+#         guidance_scale=7.5,
+#         resize_check=True,
+#         enable_safety=True,
+#         fill_mode="patchmatch",
+#         strength=0.75,
+#         step=50,
+#         enable_img2img=False,
+#         use_seed=False,
+#         seed_val=-1,
+#         generate_num=1,
+#         scheduler="",
+#         scheduler_eta=0.0,
+#         **kwargs,
+#     ):
+#         text2img, inpaint, img2img, unified = (
+#             self.text2img,
+#             self.inpaint,
+#             self.img2img,
+#             self.unified,
+#         )
+#         selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
+#         for item in [text2img, inpaint, img2img, unified]:
+#             item.scheduler = selected_scheduler
+#             if enable_safety or self.safety_checker is None:
+#                 item.safety_checker = self.safety_checker
+#             else:
+#                 item.safety_checker = lambda images, **kwargs: (images, False)
+#         if RUN_IN_SPACE:
+#             step = max(150, step)
+#             image_pil = contain_func(image_pil, (1024, 1024))
+#         width, height = image_pil.size
+#         sel_buffer = np.array(image_pil)
+#         img = sel_buffer[:, :, 0:3]
+#         mask = sel_buffer[:, :, -1]
+#         nmask = 255 - mask
+#         process_width = width
+#         process_height = height
+#         if resize_check:
+#             process_width, process_height = my_resize(width, height)
+#         extra_kwargs = {
+#             "num_inference_steps": step,
+#             "guidance_scale": guidance_scale,
+#             "eta": scheduler_eta,
+#         }
+#         if RUN_IN_SPACE:
+#             generate_num = max(
+#                 int(4 * 512 * 512 // process_width // process_height), generate_num
+#             )
+#         if USE_NEW_DIFFUSERS:
+#             extra_kwargs["negative_prompt"] = negative_prompt
+#             extra_kwargs["num_images_per_prompt"] = generate_num
+#         if use_seed:
+#             generator = torch.Generator(text2img.device).manual_seed(seed_val)
+#             extra_kwargs["generator"] = generator
+#         if nmask.sum() < 1 and enable_img2img:
+#             init_image = Image.fromarray(img)
+#             if True:
+#                 images = img2img(
+#                     prompt=prompt,
+#                     image=init_image.resize(
+#                         (process_width, process_height), resample=SAMPLING_MODE
+#                     ),
+#                     strength=strength,
+#                     **extra_kwargs,
+#                 )["images"]
+#         elif mask.sum() > 0:
+#             if fill_mode == "g_diffuser" and not self.inpainting_model:
+#                 mask = 255 - mask
+#                 mask = mask[:, :, np.newaxis].repeat(3, axis=2)
+#                 img, mask = functbl[fill_mode](img, mask)
+#                 extra_kwargs["strength"] = 1.0
+#                 extra_kwargs["out_mask"] = Image.fromarray(mask)
+#                 inpaint_func = unified
+#             else:
+#                 img, mask = functbl[fill_mode](img, mask)
+#                 mask = 255 - mask
+#                 mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+#                 mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+#                 inpaint_func = inpaint
+#             init_image = Image.fromarray(img)
+#             mask_image = Image.fromarray(mask)
+#             # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+#             input_image = init_image.resize(
+#                 (process_width, process_height), resample=SAMPLING_MODE
+#             )
+#             if self.inpainting_model:
+#                 images = inpaint_func(
+#                     prompt=prompt,
+#                     image=input_image,
+#                     width=process_width,
+#                     height=process_height,
+#                     mask_image=mask_image.resize((process_width, process_height)),
+#                     **extra_kwargs,
+#                 )["images"]
+#             else:
+#                 extra_kwargs["strength"] = strength
+#                 if True:
+#                     images = inpaint_func(
+#                         prompt=prompt,
+#                         image=input_image,
+#                         mask_image=mask_image.resize((process_width, process_height)),
+#                         **extra_kwargs,
+#                     )["images"]
+#         else:
+#             if True:
+#                 images = text2img(
+#                     prompt=prompt,
+#                     height=process_width,
+#                     width=process_height,
+#                     **extra_kwargs,
+#                 )["images"]
+#         return images
+
+
 def get_model(token="", model_choice="", model_path=""):
     if "model" not in model:
         model_name = ""