Spaces:

We-Want-GPU
/

diffusers-cross-attention-map-SDXL-t2i

Sleeping

wooyeolbaek commited on Dec 6, 2023

Commit

eeb734b

1 Parent(s): 0e62d9e

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -16,7 +16,8 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     torch_dtype=torch.float16,
 )
 pipe.unet = register_cross_attention_hook(pipe.unet)
-pipe = pipe.to("cuda")
 def inference(prompt):

     torch_dtype=torch.float16,
 )
 pipe.unet = register_cross_attention_hook(pipe.unet)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+pipe = pipe.to(device)
 def inference(prompt):