gpu

main.py

@@ -9,8 +9,8 @@ from fastapi.responses import StreamingResponse
 base = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
 )
-# base.to("cuda")
-base.enable_model_cpu_offload()
+base.to("cuda")
+# base.enable_model_cpu_offload()
 base.enable_attention_slicing()
 refiner = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-refiner-1.0",
@@ -20,8 +20,8 @@ refiner = DiffusionPipeline.from_pretrained(
     use_safetensors=True,
     variant="fp16",
 )
-# refiner.to("cuda")
-refiner.enable_model_cpu_offload()
+refiner.to("cuda")
+# refiner.enable_model_cpu_offload()
 refiner.enable_attention_slicing()
 
 # Create a new FastAPI app instance