Sophia Yang commited on
Commit
6cd7a87
1 Parent(s): 8104efd

set up device

Browse files
Files changed (1) hide show
  1. app.py +13 -3
app.py CHANGED
@@ -15,11 +15,21 @@ pn.state.template.main_max_width = "690px"
15
  pn.state.template.accent_base_color = "#F08080"
16
  pn.state.template.header_background = "#F08080"
17
 
 
 
 
18
  # Model
19
  model_id = "timbrooks/instruct-pix2pix"
20
- pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
21
- model_id, torch_dtype=torch.float16
22
- ).to("cuda")
 
 
 
 
 
 
 
23
  pipe.enable_xformers_memory_efficient_attention()
24
  pipe.unet.to(memory_format=torch.channels_last)
25
 
 
15
  pn.state.template.accent_base_color = "#F08080"
16
  pn.state.template.header_background = "#F08080"
17
 
18
+ # Set up device
19
+ device = "cuda" if torch.cuda.is_available() else "cpu"
20
+
21
  # Model
22
  model_id = "timbrooks/instruct-pix2pix"
23
+
24
+ if devide == "cuda":
25
+ pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
26
+ model_id, torch_dtype=torch.float16
27
+ )
28
+ else:
29
+ pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
30
+ model_id
31
+ )
32
+ pipe = pipe.to(device)
33
  pipe.enable_xformers_memory_efficient_attention()
34
  pipe.unet.to(memory_format=torch.channels_last)
35