update

app.py

@@ -1,5 +1,5 @@
 import random
-import spaces
+#import spaces
 
 import gradio as gr
 import numpy as np
@@ -7,8 +7,10 @@ import torch
 from diffusers import LCMScheduler, PixArtAlphaPipeline, Transformer2DModel
 from peft import PeftModel
 import os
+import devicetorch
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = devicetorch.get(torch)
+#device = "cuda" if torch.cuda.is_available() else "cpu"
 IS_SPACE = os.environ.get("SPACE_ID", None) is not None
 
 transformer = Transformer2DModel.from_pretrained(
@@ -50,7 +52,7 @@ MAX_IMAGE_SIZE = 1024
 NUM_INFERENCE_STEPS = 4
 
 
-@spaces.GPU
+#@spaces.GPU
 def infer(prompt, seed, randomize_seed):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -86,29 +88,29 @@ css = """
 }
 """
 
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
+#if torch.cuda.is_available():
+#    power_device = "GPU"
+#else:
+#    power_device = "CPU"
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(
-            f"""
-        # ⚡ Flash Diffusion: FlashPixart ⚡
-        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.*
-        [This model](https://huggingface.co/jasperai/flash-pixart) is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**. 
-        Currently running on {power_device}.
-        """
-        )
-        gr.Markdown(
-            "If you enjoy the space, please also promote *open-source* by giving a ⭐ to the <a href='https://github.com/gojasper/flash-diffusion' target='_blank'>Github Repo</a>."
-        )
-        gr.Markdown(
-            "💡 *Hint:* To better appreciate the low latency of our method, run the demo locally !"
-        )
-
+#        gr.Markdown(
+#            f"""
+#        # ⚡ Flash Diffusion: FlashPixart ⚡
+#        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.*
+#        [This model](https://huggingface.co/jasperai/flash-pixart) is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**. 
+#        Currently running on {power_device}.
+#        """
+#        )
+#        gr.Markdown(
+#            "If you enjoy the space, please also promote *open-source* by giving a ⭐ to the <a href='https://github.com/gojasper/flash-diffusion' target='_blank'>Github Repo</a>."
+#        )
+#        gr.Markdown(
+#            "💡 *Hint:* To better appreciate the low latency of our method, run the demo locally !"
+#        )
+#
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",

requirements.txt

@@ -1,13 +1,13 @@
 accelerate
 git+https://github.com/huggingface/diffusers/
 invisible_watermark
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.0.1
+#--extra-index-url https://download.pytorch.org/whl/cu118
+#torch==2.0.1
 peft >= 0.6.0
 sentencepiece==0.2.0
 optimum
 beautifulsoup4
 transformers >= 4.34.0
-xformers
+#xformers
 ftfy
-spaces
+#spaces