Spaces:

omer11a
/

bounded-attention

Runtime error

omer11a commited on Apr 4, 2024

Commit

61862e4

1 Parent(s): 333e56c

Faster loading

Files changed (1) hide show

app.py CHANGED Viewed

@@ -134,7 +134,8 @@ def inference(
     device = torch.device("cuda")
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
-    model = StableDiffusionXLPipeline.from_pretrained(LOCAL_MODEL_PATH, scheduler=scheduler, torch_dtype=torch.float16).to(device)
     model.unet.set_attn_processor(AttnProcessor2_0())
     model.enable_sequential_cpu_offload()

     device = torch.device("cuda")
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
+    model = StableDiffusionXLPipeline.from_pretrained(LOCAL_MODEL_PATH, scheduler=scheduler, torch_dtype=torch.float16, device_map="auto")
+    model.to(device)
     model.unet.set_attn_processor(AttnProcessor2_0())
     model.enable_sequential_cpu_offload()