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():