johnberg commited on
Commit
a18ea69
·
1 Parent(s): eb0e6f1

update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -5
app.py CHANGED
@@ -39,15 +39,16 @@ ckpt = torch.load(model_path, map_location='cpu')
39
  opts = ckpt['opts']
40
  opts['checkpoint_path'] = model_path
41
  opts['pretrained_e4e_path'] = e4e_path
42
- opts['device'] = 'cuda' if torch.cuda.is_available() else 'cpu'
 
43
  opts = Namespace(**opts)
44
  encoder = pSp(opts)
45
  encoder.eval()
46
- encoder.cuda()
47
 
48
  adapter = CLIPAdapterWithDecoder(opts)
49
  adapter.eval()
50
- adapter.cuda()
51
 
52
  clip_model, _ = clip.load("ViT-B/32", device='cuda')
53
 
@@ -56,8 +57,8 @@ def manipulate(input_image, caption):
56
  input_image = input_transforms(aligned_image)
57
  input_image = input_image.unsqueeze(0)
58
  text_input = clip.tokenize(caption)
59
- text_input = text_input.cuda()
60
- input_image = input_image.cuda().float()
61
 
62
  with torch.no_grad():
63
  text_features = clip_model.encode_text(text_input).float()
 
39
  opts = ckpt['opts']
40
  opts['checkpoint_path'] = model_path
41
  opts['pretrained_e4e_path'] = e4e_path
42
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
43
+ opts['device'] = device
44
  opts = Namespace(**opts)
45
  encoder = pSp(opts)
46
  encoder.eval()
47
+ encoder.to(device)
48
 
49
  adapter = CLIPAdapterWithDecoder(opts)
50
  adapter.eval()
51
+ adapter.to(device)
52
 
53
  clip_model, _ = clip.load("ViT-B/32", device='cuda')
54
 
 
57
  input_image = input_transforms(aligned_image)
58
  input_image = input_image.unsqueeze(0)
59
  text_input = clip.tokenize(caption)
60
+ text_input = text_input.to(device)
61
+ input_image = input_image.to(device).float()
62
 
63
  with torch.no_grad():
64
  text_features = clip_model.encode_text(text_input).float()