Sean Powell commited on
Commit
21ca519
·
1 Parent(s): 35483dc

Use multiple GPUs (for Nvidia A10G support).

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -9,8 +9,8 @@ img2img_strength = 0.5
9
 
10
  class EndpointHandler:
11
  def __init__(self, path=""):
12
- self.sdxl_pipe = pipes.create_stable_diffusion_xl_pipeline("cuda")
13
- self.sdxl_img2img_pipe = pipes.create_stable_diffusion_xl_img2img_pipe("cuda")
14
  print("post __init__ memory summary")
15
  print(torch.cuda.memory_summary())
16
 
 
9
 
10
  class EndpointHandler:
11
  def __init__(self, path=""):
12
+ self.sdxl_pipe = pipes.create_stable_diffusion_xl_pipeline("cuda:0")
13
+ self.sdxl_img2img_pipe = pipes.create_stable_diffusion_xl_img2img_pipe("cuda:1")
14
  print("post __init__ memory summary")
15
  print(torch.cuda.memory_summary())
16