eolecvk commited on
Commit
869347e
·
1 Parent(s): 291b964

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -7
app.py CHANGED
@@ -9,15 +9,21 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
9
  context = autocast if device == "cuda" else nullcontext
10
  dtype = torch.float16 if device == "cuda" else torch.float32
11
 
12
- if device == "cuda":
 
 
 
 
 
 
 
 
 
 
 
 
13
  pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-naruto-diffusers", torch_dtype=dtype)
14
 
15
- else:
16
- pipe = StableDiffusionOnnxPipeline.from_pretrained(
17
- "lambdalabs/sd-naruto-diffusers",
18
- revision="onnx",
19
- provider="CPUExecutionProvider"
20
- )
21
  pipe = pipe.to(device)
22
 
23
  # Sometimes the nsfw checker is confused by the Naruto images, you can disable
 
9
  context = autocast if device == "cuda" else nullcontext
10
  dtype = torch.float16 if device == "cuda" else torch.float32
11
 
12
+ try:
13
+ if device == "cuda":
14
+ pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-naruto-diffusers", torch_dtype=dtype)
15
+
16
+ else:
17
+ pipe = StableDiffusionOnnxPipeline.from_pretrained(
18
+ "lambdalabs/sd-naruto-diffusers",
19
+ revision="onnx",
20
+ provider="CPUExecutionProvider"
21
+ )
22
+
23
+ # onnx model revision not available
24
+ except:
25
  pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-naruto-diffusers", torch_dtype=dtype)
26
 
 
 
 
 
 
 
27
  pipe = pipe.to(device)
28
 
29
  # Sometimes the nsfw checker is confused by the Naruto images, you can disable