updates for CPU only diffusion

app.py

@@ -15,6 +15,7 @@ from transformers.models.detr.feature_extraction_detr import rgb_to_id
 from diffusers import StableDiffusionInpaintPipeline
 
 auth_token = os.environ.get("READ_TOKEN")
+try_cuda = True
 
 torch.inference_mode()
 torch.no_grad()
@@ -32,7 +33,7 @@ def load_diffusion_pipeline(model_name: str = 'runwayml/stable-diffusion-inpaint
     return StableDiffusionInpaintPipeline.from_pretrained(
         model_name,
         revision='fp16',
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float16 if try_cuda and torch.cuda.is_available() else torch.float32,
         use_auth_token=auth_token
     )
 
@@ -60,7 +61,7 @@ def clean_mask(mask, max_kernel: int = 23, min_kernel: int = 5):
 feature_extractor, segmentation_model, segmentation_cfg = load_segmentation_models()
 pipe = load_diffusion_pipeline()
 
-device = get_device()
+device = get_device(try_cuda=try_cuda)
 pipe = pipe.to(device)
 
 # Callback function that runs segmentation and updates CheckboxGroup
@@ -161,7 +162,9 @@ demo = gr.Blocks(css=open('app.css').read())
 
 with demo:
     gr.HTML(open('app_header.html').read())
-    # gr.Markdown("# Stable Diffusion Inpainting Segmentation")
+
+    if not try_cuda or not torch.cuda.is_available():
+        gr.HTML('<div class="alert alert-warning" role="alert" style="color:red"><b>Warning: GPU not available! Diffusion will be slow.</b></div>')
 
     # Input image control
     input_image = gr.Image(value="example.png", type='pil', label="Input Image")