Add support to run in Google Colab

app.py

@@ -2,6 +2,16 @@ import os, torch
 import gradio as gr
 from diffusers import StableDiffusionPipeline
 
+# check if running in google colab environment
+def is_google_colab():
+    try:
+        import google.colab
+        return True
+    except:
+        return False
+
+is_colab = is_google_colab()
+
 # get hf user access token as an environment variable
 TOKEN_KEY = os.getenv('AUTH_TOKEN')
 
@@ -23,9 +33,9 @@ def generate(prompt:str, seed:int, guidance:float, steps:int):
 
 if device == "cuda":
     device_name = torch.cuda.get_device_name(0)
-    print(device_name + " available.")
+    print(device_name + " available! Using GPU 🔥")
 else:
-    print("Using CPU.")
+    print("Using CPU😒.")
 
 # create the gradio UI
 # set precision to 0 to round value to nearest int
@@ -40,5 +50,4 @@ demo = gr.Interface(
 demo.queue(concurrency_count=3)
 
 # launch demo
-# demo.launch(share=True) # uncomment if running outside of huggingface spaces
-demo.launch()
+demo.launch(debug=is_colab, share=is_colab)