Add google colab

app.py

@@ -3,18 +3,27 @@ import gradio as gr
 import torch
 from torch import autocast
 from diffusers import DiffusionPipeline
-import streamlit as st
 from transformers import (
     pipeline,
     MBart50TokenizerFast,
     MBartForConditionalGeneration,
 )
 
+import utils
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 device_dict = {"cuda": 0, "cpu": -1}
 context = autocast if device == "cuda" else nullcontext
 dtype = torch.float16 if device == "cuda" else torch.float32
 
+# Detect if code is running in Colab
+is_colab = utils.is_google_colab()
+colab_instruction = "" if is_colab else """
+<p>You can skip the queue using Colab: 
+<a href="https://colab.research.google.com/drive/1nhXyddThldnxPfIYO2my_bYinlMUW30R?usp=sharing">
+<img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a></p>"""
+device_print = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
+
 # Add language detection pipeline
 language_detection_model_ckpt = "papluca/xlm-roberta-base-language-detection"
 language_detection_pipeline = pipeline("text-classification",
@@ -27,16 +36,28 @@ trans_model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-larg
 
 model_id = "CompVis/stable-diffusion-v1-4"
 
-pipe = DiffusionPipeline.from_pretrained(
-    model_id,
-    custom_pipeline="multilingual_stable_diffusion",
-    use_auth_token=st.secrets["USER_TOKEN"],
-    detection_pipeline=language_detection_pipeline,
-    translation_model=trans_model,
-    translation_tokenizer=trans_tokenizer,
-    revision="fp16",
-    torch_dtype=dtype,
-)
+if is_colab:
+    pipe = DiffusionPipeline.from_pretrained(
+        model_id,
+        custom_pipeline="multilingual_stable_diffusion",
+        detection_pipeline=language_detection_pipeline,
+        translation_model=trans_model,
+        translation_tokenizer=trans_tokenizer,
+        revision="fp16",
+        torch_dtype=dtype,
+    )
+else:
+    import streamlit as st
+    pipe = DiffusionPipeline.from_pretrained(
+        model_id,
+        custom_pipeline="multilingual_stable_diffusion",
+        use_auth_token=st.secrets["USER_TOKEN"],
+        detection_pipeline=language_detection_pipeline,
+        translation_model=trans_model,
+        translation_tokenizer=trans_tokenizer,
+        revision="fp16",
+        torch_dtype=dtype,
+    )
 
 pipe = pipe.to(device)
 
@@ -228,6 +249,10 @@ with block as demo:
               <p style="margin-bottom: 10px; font-size: 94%">
                 Stable Diffusion Pipeline that supports prompts in 50 different languages.
               </p>
+              <p>
+                {colab_instruction}
+                Running on <b>{device_print}</b>{(" in a <b>Google Colab</b>." if is_colab else "")}
+              </p>
             </div>
         """
     )
@@ -277,4 +302,8 @@ Despite how impressive being able to turn text into image is, beware to the fact
                </div>
            """
         )
+    gr.Markdown('''
+      [![Twitter Follow](https://img.shields.io/twitter/follow/juancopi81?style=social)](https://twitter.com/juancopi81)
+      ![visitors](https://visitor-badge.glitch.me/badge?page_id=Juancopi81.MultilingualStableDiffusion)
+    ''')
 demo.queue(max_size=25).launch()

utils.py

@@ -0,0 +1,6 @@
+def is_google_colab():
+    try:
+        import google.colab
+        return True
+    except:
+        return False