Update app.py

app.py

@@ -78,25 +78,24 @@ def generate(
     use_resolution_binning: bool = True,
 ):
     
-    if torch.cuda.is_available():
-        pipe = DiffusionPipeline.from_pretrained(
-            "playgroundai/playground-v2.5-1024px-aesthetic",
-            torch_dtype=torch.float16,
-            use_safetensors=True,
-            add_watermarker=False,
-            variant="fp16"
-        )
-        if ENABLE_CPU_OFFLOAD:
-            pipe.enable_model_cpu_offload()
-        else:
-            pipe.to(device)    
-            print("Loaded on Device!")
-        
-        if USE_TORCH_COMPILE:
-            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-            print("Model Compiled!")
+    pipe = DiffusionPipeline.from_pretrained(
+        "playgroundai/playground-v2.5-1024px-aesthetic",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        add_watermarker=False,
+        variant="fp16"
+    )
+    if ENABLE_CPU_OFFLOAD:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe.to(device)    
+        print("Loaded on Device!")
+    
+    if USE_TORCH_COMPILE:
+        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        print("Model Compiled!")
 
-    pipe.to(device)
+    pipe.to('cpu')
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)