nightfury commited on
Commit
02df5d4
1 Parent(s): 6c4be6b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -3
app.py CHANGED
@@ -21,15 +21,20 @@ def download_image(url):
21
  response = requests.get(url)
22
  return PIL.Image.open(BytesIO(response.content)).convert("RGB")
23
 
24
- device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 
25
  pipe = StableDiffusionInpaintingPipeline.from_pretrained(
26
  "CompVis/stable-diffusion-v1-4",
27
- revision="fp16",
28
  torch_dtype=torch.float16,
29
  use_auth_token=auth_token,
30
  ).to(device)
31
 
32
- model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
 
 
33
  model.eval()
34
  model.load_state_dict(torch.load('./clipseg/weights/rd64-uni.pth', map_location=torch.device('cuda')), strict=False)
35
 
21
  response = requests.get(url)
22
  return PIL.Image.open(BytesIO(response.content)).convert("RGB")
23
 
24
+ #device = "cuda" if torch.cuda.is_available() else "cpu"
25
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
26
+ print("The model will be running on :: ", device, " ~device")
27
+
28
  pipe = StableDiffusionInpaintingPipeline.from_pretrained(
29
  "CompVis/stable-diffusion-v1-4",
30
+ #revision="fp16",
31
  torch_dtype=torch.float16,
32
  use_auth_token=auth_token,
33
  ).to(device)
34
 
35
+ #model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
36
+ model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True)
37
+
38
  model.eval()
39
  model.load_state_dict(torch.load('./clipseg/weights/rd64-uni.pth', map_location=torch.device('cuda')), strict=False)
40