eyal.benaroche commited on
Commit
dc5d4c5
1 Parent(s): e7af588

update app readme + disclaimer

Browse files
Files changed (1) hide show
  1. app.py +50 -53
app.py CHANGED
@@ -1,38 +1,35 @@
 
 
1
  import gradio as gr
2
  import numpy as np
3
- import random
4
- from diffusers import PixArtAlphaPipeline, Transformer2DModel, LCMScheduler
5
  import torch
 
6
  from peft import PeftModel
7
 
8
-
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
 
11
  transformer = Transformer2DModel.from_pretrained(
12
- "PixArt-alpha/PixArt-XL-2-1024-MS",
13
- subfolder="transformer",
14
- torch_dtype=torch.float16
15
- )
16
- transformer = PeftModel.from_pretrained(
17
- transformer,
18
- "jasperai/flash-pixart"
19
  )
 
20
 
21
 
22
  if torch.cuda.is_available():
23
  torch.cuda.max_memory_allocated(device=device)
24
  pipe = PixArtAlphaPipeline.from_pretrained(
25
- "PixArt-alpha/PixArt-XL-2-1024-MS",
26
- transformer=transformer,
27
- torch_dtype=torch.float16
28
  )
29
  pipe.enable_xformers_memory_efficient_attention()
30
  pipe = pipe.to(device)
31
- else:
32
  pipe = PixArtAlphaPipeline.from_pretrained(
33
- "PixArt-alpha/PixArt-XL-2-1024-MS",
34
- transformer=transformer,
35
- torch_dtype=torch.float16
36
  )
37
  pipe = pipe.to(device)
38
 
@@ -48,22 +45,24 @@ MAX_SEED = np.iinfo(np.int32).max
48
  MAX_IMAGE_SIZE = 1024
49
  NUM_INFERENCE_STEPS = 4
50
 
 
51
  def infer(prompt, seed, randomize_seed):
52
 
53
  if randomize_seed:
54
  seed = random.randint(0, MAX_SEED)
55
-
56
  generator = torch.Generator().manual_seed(seed)
57
-
58
  image = pipe(
59
- prompt = prompt,
60
- guidance_scale = 0,
61
- num_inference_steps = NUM_INFERENCE_STEPS,
62
- generator = generator
63
- ).images[0]
64
-
65
  return image
66
 
 
67
  examples = [
68
  "The image showcases a freshly baked bread, possibly focaccia, with rosemary sprigs and red pepper flakes sprinkled on top. It's sliced and placed on a wire cooling rack, with a bowl of mixed peppercorns beside it.",
69
  "A raccoon reading a book in a lush forest.",
@@ -76,7 +75,7 @@ examples = [
76
  "A beautiful sunflower in rainy day",
77
  ]
78
 
79
- css="""
80
  #col-container {
81
  margin: 0 auto;
82
  max-width: 512px;
@@ -89,17 +88,19 @@ else:
89
  power_device = "CPU"
90
 
91
  with gr.Blocks(css=css) as demo:
92
-
93
  with gr.Column(elem_id="col-container"):
94
- gr.Markdown(f"""
 
95
  # ⚡ FlashDiffusion: FlashPixart ⚡
96
- This is an interactive demo of [Flash Diffusion](https://huggingface.co/jasperai/flash-pixart), a diffusion distillation method proposed in [ADD ARXIV]() *by Clément Chadebec, Onur Tasar and Benjamin Aubin.*
97
- This model 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**.
98
  Currently running on {power_device}.
99
- """)
100
-
 
101
  with gr.Row():
102
-
103
  prompt = gr.Text(
104
  label="Prompt",
105
  show_label=False,
@@ -107,13 +108,13 @@ with gr.Blocks(css=css) as demo:
107
  placeholder="Enter your prompt",
108
  container=False,
109
  )
110
-
111
  run_button = gr.Button("Run", scale=0)
112
-
113
  result = gr.Image(label="Result", show_label=False)
114
 
115
  with gr.Accordion("Advanced Settings", open=False):
116
-
117
  seed = gr.Slider(
118
  label="Seed",
119
  minimum=0,
@@ -121,24 +122,20 @@ with gr.Blocks(css=css) as demo:
121
  step=1,
122
  value=0,
123
  )
124
-
125
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
126
-
127
-
128
- gr.Examples(
129
- examples = examples,
130
- inputs = [prompt]
131
  )
132
 
133
- run_button.click(
134
- fn = infer,
135
- inputs = [prompt, seed, randomize_seed],
136
- outputs = [result]
137
- )
138
- seed.change(
139
- fn = infer,
140
- inputs = [prompt, seed, randomize_seed],
141
- outputs = [result]
142
- )
143
 
144
- demo.queue().launch()
 
1
+ import random
2
+
3
  import gradio as gr
4
  import numpy as np
 
 
5
  import torch
6
+ from diffusers import LCMScheduler, PixArtAlphaPipeline, Transformer2DModel
7
  from peft import PeftModel
8
 
 
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
 
11
  transformer = Transformer2DModel.from_pretrained(
12
+ "PixArt-alpha/PixArt-XL-2-1024-MS",
13
+ subfolder="transformer",
14
+ torch_dtype=torch.float16,
 
 
 
 
15
  )
16
+ transformer = PeftModel.from_pretrained(transformer, "jasperai/flash-pixart")
17
 
18
 
19
  if torch.cuda.is_available():
20
  torch.cuda.max_memory_allocated(device=device)
21
  pipe = PixArtAlphaPipeline.from_pretrained(
22
+ "PixArt-alpha/PixArt-XL-2-1024-MS",
23
+ transformer=transformer,
24
+ torch_dtype=torch.float16,
25
  )
26
  pipe.enable_xformers_memory_efficient_attention()
27
  pipe = pipe.to(device)
28
+ else:
29
  pipe = PixArtAlphaPipeline.from_pretrained(
30
+ "PixArt-alpha/PixArt-XL-2-1024-MS",
31
+ transformer=transformer,
32
+ torch_dtype=torch.float16,
33
  )
34
  pipe = pipe.to(device)
35
 
 
45
  MAX_IMAGE_SIZE = 1024
46
  NUM_INFERENCE_STEPS = 4
47
 
48
+
49
  def infer(prompt, seed, randomize_seed):
50
 
51
  if randomize_seed:
52
  seed = random.randint(0, MAX_SEED)
53
+
54
  generator = torch.Generator().manual_seed(seed)
55
+
56
  image = pipe(
57
+ prompt=prompt,
58
+ guidance_scale=0,
59
+ num_inference_steps=NUM_INFERENCE_STEPS,
60
+ generator=generator,
61
+ ).images[0]
62
+
63
  return image
64
 
65
+
66
  examples = [
67
  "The image showcases a freshly baked bread, possibly focaccia, with rosemary sprigs and red pepper flakes sprinkled on top. It's sliced and placed on a wire cooling rack, with a bowl of mixed peppercorns beside it.",
68
  "A raccoon reading a book in a lush forest.",
 
75
  "A beautiful sunflower in rainy day",
76
  ]
77
 
78
+ css = """
79
  #col-container {
80
  margin: 0 auto;
81
  max-width: 512px;
 
88
  power_device = "CPU"
89
 
90
  with gr.Blocks(css=css) as demo:
91
+
92
  with gr.Column(elem_id="col-container"):
93
+ gr.Markdown(
94
+ f"""
95
  # ⚡ FlashDiffusion: FlashPixart ⚡
96
+ This is an interactive demo of [Flash Diffusion](https://flash-diffusion.gojasper.github.io), a diffusion distillation method proposed in [ADD ARXIV]() *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin.*
97
+ [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**.
98
  Currently running on {power_device}.
99
+ """
100
+ )
101
+
102
  with gr.Row():
103
+
104
  prompt = gr.Text(
105
  label="Prompt",
106
  show_label=False,
 
108
  placeholder="Enter your prompt",
109
  container=False,
110
  )
111
+
112
  run_button = gr.Button("Run", scale=0)
113
+
114
  result = gr.Image(label="Result", show_label=False)
115
 
116
  with gr.Accordion("Advanced Settings", open=False):
117
+
118
  seed = gr.Slider(
119
  label="Seed",
120
  minimum=0,
 
122
  step=1,
123
  value=0,
124
  )
125
+
126
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
127
+
128
+ examples = gr.Examples(examples=examples, inputs=[prompt])
129
+
130
+ gr.Markdown(
131
+ "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."
132
  )
133
 
134
+ gr.Markdown(
135
+ "To better appreciate the low latency of our method, run the demo locally !"
136
+ )
137
+
138
+ run_button.click(fn=infer, inputs=[prompt, seed, randomize_seed], outputs=[result])
139
+ seed.change(fn=infer, inputs=[prompt, seed, randomize_seed], outputs=[result])
 
 
 
 
140
 
141
+ demo.queue().launch(show_api=False)