taesd3

README.md

@@ -4,10 +4,10 @@ emoji: ⚡
 colorFrom: green
 colorTo: blue
 sdk: gradio
-sdk_version: 4.26.0
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 license: cc-by-nc-4.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -4,7 +4,12 @@ import spaces
 import gradio as gr
 import numpy as np
 import torch
-from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlashFlowMatchEulerDiscreteScheduler
+from diffusers import (
+    StableDiffusion3Pipeline,
+    SD3Transformer2DModel,
+    FlashFlowMatchEulerDiscreteScheduler,
+    AutoencoderTiny,
+)
 from peft import PeftModel
 import os
 from huggingface_hub import snapshot_download
@@ -12,16 +17,15 @@ from huggingface_hub import snapshot_download
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
 model_path = snapshot_download(
-    repo_id="stabilityai/stable-diffusion-3-medium", 
-    revision="refs/pr/26",
-    repo_type="model", 
+    repo_id="stabilityai/stable-diffusion-3-medium-diffusers",
+    repo_type="model",
     ignore_patterns=["*.md", "*..gitattributes"],
     local_dir="stable-diffusion-3-medium",
-    token=huggingface_token, # type a new token-id.
-    )
+    token=huggingface_token,  # type a new token-id.
+)
+import spaces
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-IS_SPACE = os.environ.get("SPACE_ID", None) is not None
 
 transformer = SD3Transformer2DModel.from_pretrained(
     model_path,
@@ -40,6 +44,10 @@ if torch.cuda.is_available():
         text_encoder_3=None,
         tokenizer_3=None,
     )
+    pipe.vae = AutoencoderTiny.from_pretrained(
+        "madebyollin/taesd3", torch_dtype=torch.float16
+    )
+    pipe.vae.config.shift_factor = 0.0
 
     pipe = pipe.to(device)
 else:
@@ -54,9 +62,11 @@ else:
 
 
 pipe.scheduler = FlashFlowMatchEulerDiscreteScheduler.from_pretrained(
-  model_path,
-  subfolder="scheduler",
+    model_path,
+    subfolder="scheduler",
 )
+pipe.set_progress_bar_config(disable=True)
+
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -89,7 +99,7 @@ examples = [
     "a cute cartoon fluffy rabbit pilot walking on a military aircraft carrier, 8k, cinematic",
     "A 3d render of a futuristic city with a giant robot in the middle full of neon lights, pink and blue colors",
     "A close up of an old elderly man with green eyes looking straight at the camera",
-    "photo of a huge red cat with green eyes sitting on a cloud in the sky, looking at the camera"
+    "photo of a huge red cat with green eyes sitting on a cloud in the sky, looking at the camera",
 ]
 
 css = """
@@ -108,12 +118,8 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(
             f"""
-        # ⚡ Flash Diffusion: FlashSD3 ⚡
-        This is an interactive demo of [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/), a diffusion distillation method proposed in [Flash Diffusion: Accelerating Any Conditional
-        Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin* from Jasper Research.
-        [This model](https://huggingface.co/jasperai/flash-sd3) is a **90.4M** LoRA distilled version of [SD3](https://huggingface.co/stabilityai/stable-diffusion-3-medium) model that is able to generate 1024x1024 images in **4 steps**.
-        Results can be compared with the teacher model [here](https://huggingface.co/spaces/stabilityai/stable-diffusion-3-medium).
-        Currently running on {power_device}.
+        # ⚡ Flash Diffusion: FlashSD3 + TAESD3 ⚡️
+        [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/) with [Tiny AutoEncoder for Stable Diffusion 3](https://huggingface.co/madebyollin/taesd3)
         """
         )
         gr.Markdown(
@@ -154,7 +160,13 @@ with gr.Blocks(css=css) as demo:
             "This demo is only for research purpose. Jasper cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. Jasper provides the tools, but the responsibility for their use lies with the individual user."
         )
     gr.on(
-        [run_button.click, seed.change, randomize_seed.change, prompt.submit],
+        [
+            run_button.click,
+            seed.change,
+            randomize_seed.change,
+            prompt.submit,
+            prompt.change,
+        ],
         fn=infer,
         inputs=[prompt, seed, randomize_seed],
         outputs=[result],

requirements.txt

@@ -1,14 +1,15 @@
 accelerate
 diffusers @ git+https://github.com/initml/diffusers.git@clement/feature/flash_sd3
 invisible_watermark
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.0.1
+--extra-index-url https://download.pytorch.org/whl/cu221
+torch
 transformers >= 4.34.0
-xformers
+# xformers
 ftfy
 spaces
 peft >= 0.6.0
 sentencepiece==0.2.0
 optimum
 beautifulsoup4
-gradio==4.29.0
+gradio==4.36.1
+setuptools