ChenWu98 patrickvonplaten commited on
Commit
2bff877
·
1 Parent(s): 04d4d07

Update app.py (#2)

Browse files

- Update app.py (a0cce46c41ec54b3151b8ea82f117d6f85be4163)


Co-authored-by: Patrick von Platen <[email protected]>

Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -23,7 +23,8 @@ if is_colab:
23
  else:
24
  import streamlit as st
25
  scheduler = DDIMScheduler.from_config(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], subfolder="scheduler")
26
- pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], scheduler=scheduler)
 
27
  tokenizer = pipe.tokenizer
28
 
29
  if torch.cuda.is_available():
 
23
  else:
24
  import streamlit as st
25
  scheduler = DDIMScheduler.from_config(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], subfolder="scheduler")
26
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
27
+ pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], scheduler=scheduler, torch_dtype=torch_dtype)
28
  tokenizer = pipe.tokenizer
29
 
30
  if torch.cuda.is_available():