Hugging Face's logo

Spaces:
AnTo2209
/
3D_Zeroshot_Neural_Style_Transfer
Runtime error

App Files Community
AnTo2209 commited on
Commit
da29c37
·
1 Parent(s): 594cdd5

add inference

Browse files
Files changed (1) hide show
  1. inference.py +4 -1
inference.py CHANGED
@@ -32,7 +32,10 @@ def inference(cfg, render_mode: str, image=None):
32
 
33
  trans = T.Compose([T.Resize(size=(256, 256)), T.ToTensor()])
34
  if image:
35
- style_img = trans(image).cuda()[None, ...]
 
 
 
36
  else:
37
  style_img = trans(Image.open(cfg["global"]["style_img"])).cuda()[None, ...]
38
  style_name = Path(cfg["global"]["style_img"]).stem
 
32
 
33
  trans = T.Compose([T.Resize(size=(256, 256)), T.ToTensor()])
34
  if image:
35
+ if torch.cuda.is_available():
36
+ style_img = trans(image).cuda()[None, ...]
37
+ else:
38
+ style_img = trans(image)[None, ...]
39
  else:
40
  style_img = trans(Image.open(cfg["global"]["style_img"])).cuda()[None, ...]
41
  style_name = Path(cfg["global"]["style_img"]).stem