Update app.py

app.py

@@ -3,7 +3,6 @@ import gradio as gr
 import torch
 from PIL import Image
 import utils
-import streamlit as st
 import ptp_utils
 import seq_aligner
 import torch.nn.functional as nnf
@@ -16,18 +15,18 @@ MAX_NUM_WORDS = 77
 is_colab = utils.is_google_colab()
 
 
-if True:
-    model_id_or_path = "CompVis/stable-diffusion-v1-4"
-    scheduler = DDIMScheduler.from_config(model_id_or_path,
-                                          use_auth_token=st.secrets["USER_TOKEN"],
-                                          subfolder="scheduler")
-    pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path,
-                                                  use_auth_token=st.secrets["USER_TOKEN"],
-                                                  scheduler=scheduler)
-    tokenizer = pipe.tokenizer
+model_id_or_path = "CompVis/stable-diffusion-v1-4"
+if is_colab:
+    scheduler = DDIMScheduler.from_config(model_id_or_path, subfolder="scheduler")
+    pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, scheduler=scheduler)
+else:
+    import streamlit as st
+    scheduler = DDIMScheduler.from_config(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], subfolder="scheduler")
+    pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, use_auth_token=st.secrets["USER_TOKEN"], scheduler=scheduler)
+tokenizer = pipe.tokenizer
 
-    if torch.cuda.is_available():
-        pipe = pipe.to("cuda")
+if torch.cuda.is_available():
+    pipe = pipe.to("cuda")
 
 device_print = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 device = "cuda" if torch.cuda.is_available() else "cpu"