Update src/pipeline.py

src/pipeline.py

@@ -19,11 +19,13 @@ def callback_dynamic_cfg(pipe, step_index, timestep, callback_kwargs):
   return callback_kwargs
 
 def load_pipeline(pipeline=None) -> StableDiffusionXLPipeline:
+    device = 'cpu'
     if not pipeline:
+        device = "cuda"
         pipeline = StableDiffusionXLPipeline.from_pretrained(
             "stablediffusionapi/newdream-sdxl-20",
             torch_dtype=torch.float16,
-        ).to("cuda")
+        ).to(device)
     
     pipeline.scheduler = SchedulerWrapper(DDIMScheduler.from_config(pipeline.scheduler.config))