estellea commited on
Commit
c49ce5c
1 Parent(s): 942501f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -11,9 +11,10 @@ Path("tmp").mkdir(exist_ok=True)
11
 
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
  print(f"Device is {device}")
 
14
  pipe = StableDiffusionLDM3DPipeline.from_pretrained(
15
  "Intel/ldm3d-4c",
16
- torch_dtype=torch.float16
17
  # , safety_checker=None
18
  )
19
  pipe.to(device)
 
11
 
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
  print(f"Device is {device}")
14
+ torch_type = torch.float16 if device == "cuda" else torch.float32
15
  pipe = StableDiffusionLDM3DPipeline.from_pretrained(
16
  "Intel/ldm3d-4c",
17
+ torch_dtype=torch_type
18
  # , safety_checker=None
19
  )
20
  pipe.to(device)