omer11a commited on
Commit
61862e4
·
1 Parent(s): 333e56c

Faster loading

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -134,7 +134,8 @@ def inference(
134
 
135
  device = torch.device("cuda")
136
  scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
137
- model = StableDiffusionXLPipeline.from_pretrained(LOCAL_MODEL_PATH, scheduler=scheduler, torch_dtype=torch.float16).to(device)
 
138
  model.unet.set_attn_processor(AttnProcessor2_0())
139
  model.enable_sequential_cpu_offload()
140
 
 
134
 
135
  device = torch.device("cuda")
136
  scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
137
+ model = StableDiffusionXLPipeline.from_pretrained(LOCAL_MODEL_PATH, scheduler=scheduler, torch_dtype=torch.float16, device_map="auto")
138
+ model.to(device)
139
  model.unet.set_attn_processor(AttnProcessor2_0())
140
  model.enable_sequential_cpu_offload()
141