adds stable diffusion 2, attention slicing, cuda masking

app.py

@@ -6,16 +6,24 @@ import os
 from PIL import Image
 from typing import List, Optional
 from functools import reduce
+from argparse import ArgumentParser
 
 import gradio as gr
 
 from transformers import DetrFeatureExtractor, DetrForSegmentation, DetrConfig
 from transformers.models.detr.feature_extraction_detr import rgb_to_id
 
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInpaintPipeline, EulerDiscreteScheduler
+
+# TODO: xformers install for faster diffusion
+
+parser = ArgumentParser()
+parser.add_argument('--disable-cuda', action='store_true')
+parser.add_argument('--attention-slicing', action='store_true')
+args = parser.parse_args()
 
 auth_token = os.environ.get("READ_TOKEN")
-try_cuda = True
+try_cuda = not args.disable_cuda
 
 torch.inference_mode()
 torch.no_grad()
@@ -29,7 +37,7 @@ def load_segmentation_models(model_name: str = 'facebook/detr-resnet-50-panoptic
     return feature_extractor, model, cfg
 
 # Load diffusion pipeline
-def load_diffusion_pipeline(model_name: str = 'runwayml/stable-diffusion-inpainting'):
+def load_diffusion_pipeline(model_name: str = 'stabilityai/stable-diffusion-2-inpainting'):
     return StableDiffusionInpaintPipeline.from_pretrained(
         model_name,
         revision='fp16',
@@ -51,10 +59,10 @@ def max_pool(x: torch.Tensor, kernel_size: int):
 
 # Apply min-max pooling to clean up mask
 def clean_mask(mask, max_kernel: int = 23, min_kernel: int = 5):
-    mask = torch.Tensor(mask[None, None]).float()
+    mask = torch.Tensor(mask[None, None]).float().to(device)
     mask = min_pool(mask, min_kernel)
     mask = max_pool(mask, max_kernel)
-    mask = mask.bool().squeeze().numpy()
+    mask = mask.bool().squeeze().cpu().numpy()
     return mask
 
 
@@ -62,11 +70,14 @@ feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_mode
 pipe = load_diffusion_pipeline()
 
 device = get_device(try_cuda=try_cuda)
+segmentation_model = segmentation_model.to(device)
 pipe = pipe.to(device)
+if args.attention_slicing:
+    pipe.enable_attention_slicing()
 
 # Callback function that runs segmentation and updates CheckboxGroup
 def fn_segmentation(image, max_kernel, min_kernel):
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs = feature_extractor(images=image, return_tensors="pt").to(device)
     outputs = segmentation_model(**inputs)
 
     processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)