Update app.py

app.py

@@ -4,6 +4,8 @@ import random
 import spaces
 import torch
 from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
+from PIL import Image
+import io
 
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -28,6 +30,12 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidan
     ).images[0]
     return image, seed
 
+def download_image(image, file_format):
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format=file_format)
+    img_byte_arr = img_byte_arr.getvalue()
+    return img_byte_arr
+
 examples = [
     "a galaxy swirling with vibrant blue and purple hues",
     "a futuristic cityscape under a dark sky",
@@ -153,6 +161,21 @@ with gr.Blocks(css=css) as demo:
                     value=28,
                 )
         
+        download_format = gr.Radio(
+            label="Download Format",
+            choices=["PNG", "JPEG", "SVG"],
+            value="PNG",
+            type="value",
+        )
+
+        download_button = gr.Button("Download Image")
+
+        download_button.click(
+            fn=download_image,
+            inputs=[result, download_format],
+            outputs=gr.File(label="Download"),
+        )
+
         gr.Examples(
             examples=examples,
             fn=infer,
@@ -161,6 +184,14 @@ with gr.Blocks(css=css) as demo:
             cache_examples="lazy"
         )
 
+        share_button = gr.Button("Share on Twitter")
+        share_button.click(
+            fn=lambda img: f"https://twitter.com/intent/post?url={img}",
+            inputs=[result],
+            outputs=None,
+            _js="(img) => window.open(img, '_blank')"
+        )
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,