Update app/main.py

app/main.py

@@ -39,7 +39,8 @@ async def lifespan(app: FastAPI):
     ti_pipe.scheduler = DPMSolverMultistepScheduler.from_config(
         ti_pipe.scheduler.config
     )
-    ti_pipe.enable_sequential_cpu_offload()  # to(DEVICE)
+    ti_pipe.to(DEVICE)
+    ti_pipe.enable_xformers_memory_efficient_attention()
 
     ii_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
         "timbrooks/instruct-pix2pix", torch_dtype=torch.float16, safety_checker=None
@@ -47,7 +48,8 @@ async def lifespan(app: FastAPI):
     ii_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
         ii_pipe.scheduler.config
     )
-    ii_pipe.enable_sequential_cpu_offload()  # to(DEVICE)
+    ii_pipe.to(DEVICE)
+    ii_pipe.enable_xformers_memory_efficient_attention()
 
     br_model = AutoModelForImageSegmentation.from_pretrained(
         "briaai/RMBG-1.4", trust_remote_code=True