Update app.py

cuda -> cpu

app.py

@@ -20,10 +20,10 @@ checkpoints = {
 }
 loaded = None
 
-
+pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16")
 # Ensure model and scheduler are initialized in GPU-enabled function
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
+#if torch.cuda.is_available():
+    #pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
 
 if SAFETY_CHECKER:
     from safety_checker import StableDiffusionSafetyChecker
@@ -39,10 +39,12 @@ if SAFETY_CHECKER:
     def check_nsfw_images(
         images: list[Image.Image],
     ) -> tuple[list[Image.Image], list[bool]]:
-        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
+        #safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
+        safety_checker_input = feature_extractor(images, return_tensors="pt")
         has_nsfw_concepts = safety_checker(
             images=[images],
-            clip_input=safety_checker_input.pixel_values.to("cuda")
+            #clip_input=safety_checker_input.pixel_values.to("cuda")
+            clip_input=safety_checker_input.pixel_values
         )
 
         return images, has_nsfw_concepts
@@ -58,7 +60,8 @@ def generate_image(prompt, ckpt):
 
     if loaded != num_inference_steps:
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample" if num_inference_steps==1 else "epsilon")
-        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
+        #pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
+        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cpu"))
         loaded = num_inference_steps
         
     results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)