Update stable_cascade.py

stable_cascade.py

@@ -1,12 +1,37 @@
 import torch, os
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 import gradio as gr
+from io import BytesIO
+import base64
+import re
+
+SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
+
+# Regex pattern to match data URI scheme
+data_uri_pattern = re.compile(r'data:image/(png|jpeg|jpg|webp);base64,')
+
+def readb64(b64):
+    # Remove any data URI scheme prefix with regex
+    b64 = data_uri_pattern.sub("", b64)
+    # Decode and open the image with PIL
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    return img
+    
+# convert from PIL to base64
+def writeb64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    b64image = base64.b64encode(buffered.getvalue())
+    b64image_str = b64image.decode("utf-8")
+    return b64image_str
+
 
 prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to("cuda")
 decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to("cuda")
 
 def generate_images(
-    prompt="a photo of a girl",
+    secret_token="",
+    prompt="",
     negative_prompt="bad,ugly,deformed",
     height=1024,
     width=1024, 
@@ -29,6 +54,10 @@ def generate_images(
     Returns:
     - List[PIL.Image]: A list of generated PIL Image objects.
     """
+    if secret_token != SECRET_TOKEN:
+        raise gr.Error(
+            f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
+
     generator = torch.Generator(device="cuda").manual_seed(int(seed))
 
     # Generate image embeddings using the prior model
@@ -54,13 +83,21 @@ def generate_images(
         num_inference_steps=decoder_inference_steps
     ).images
 
-    return decoder_output
+    image = decoder_output[0]
+        
+    image_base64 = writeb64(image)
+        
+    return image_base64
 
 
 def web_demo():
     with gr.Blocks():
         with gr.Row():
             with gr.Column():
+                secret_token = gr.Textbox(
+                    placeholder="Secret token",
+                    show_label=False,
+                )
                 text2image_prompt = gr.Textbox(
                     lines=1,
                     placeholder="Prompt",
@@ -129,16 +166,12 @@ def web_demo():
                                 )               
                 text2image_predict = gr.Button(value="Generate Image")
                 
-            with gr.Column():
-                output_image = gr.Gallery(
-                    label="Generated images",
-                    show_label=False,
-                    elem_id="gallery",
-                ).style(grid=(1, 2), height=300)
+           output_image_base64 = gr.Text()
                 
             text2image_predict.click(
                 fn=generate_images,
                 inputs=[
+                    secret_token,
                     text2image_prompt,
                     text2image_negative_prompt,
                     text2image_height,
@@ -149,5 +182,5 @@ def web_demo():
                     text2image_prior_inference_step,
                     text2image_decoder_inference_step
                 ],
-                outputs=output_image,
+                outputs=output_image_base64,
             )