Sophia Yang commited on
Commit
0faa1c8
1 Parent(s): 3d73a5d

switch between gpu and cpu

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -17,13 +17,14 @@ pn.state.template.param.update(
17
  )
18
 
19
  model_id = "timbrooks/instruct-pix2pix"
 
20
 
21
  if 'pipe' in pn.state.cache:
22
  pipe = pn.state.cache['pipe']
23
  else:
24
  pipe = pn.state.cache['pipe'] = StableDiffusionInstructPix2PixPipeline.from_pretrained(
25
  model_id, torch_dtype=torch.float16
26
- ).to("cuda")
27
  pipe.enable_xformers_memory_efficient_attention()
28
  pipe.unet.to(memory_format=torch.channels_last)
29
 
@@ -38,6 +39,9 @@ def normalize_image(value, width):
38
  return image.resize((width, height), PIL.Image.LANCZOS)
39
 
40
  def new_image(prompt, image, img_guidance, guidance, steps, width=600):
 
 
 
41
  edit = pipe(
42
  prompt,
43
  image=image,
 
17
  )
18
 
19
  model_id = "timbrooks/instruct-pix2pix"
20
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
21
 
22
  if 'pipe' in pn.state.cache:
23
  pipe = pn.state.cache['pipe']
24
  else:
25
  pipe = pn.state.cache['pipe'] = StableDiffusionInstructPix2PixPipeline.from_pretrained(
26
  model_id, torch_dtype=torch.float16
27
+ ).to(device)
28
  pipe.enable_xformers_memory_efficient_attention()
29
  pipe.unet.to(memory_format=torch.channels_last)
30
 
 
39
  return image.resize((width, height), PIL.Image.LANCZOS)
40
 
41
  def new_image(prompt, image, img_guidance, guidance, steps, width=600):
42
+ """
43
+ create a new image from the StableDiffusionInstructPix2PixPipeline model
44
+ """
45
  edit = pipe(
46
  prompt,
47
  image=image,