sayakpaul HF staff commited on
Commit
7c36275
1 Parent(s): 5dab465

modify space

Browse files
Files changed (2) hide show
  1. app.py +49 -277
  2. requirements.txt +4 -7
app.py CHANGED
@@ -1,214 +1,80 @@
1
- #!/usr/bin/env python
2
-
3
- from __future__ import annotations
4
-
5
  import os
6
  import random
7
 
8
  import gradio as gr
9
  import numpy as np
10
- import PIL.Image
11
  import torch
12
- from diffusers import AutoencoderKL, StableDiffusionXLPipeline
 
13
  import uuid
14
 
15
- DESCRIPTION = '''# Segmind Stable Diffusion: SSD-1B
16
- #### [Segmind's SSD-1B](https://huggingface.co/segmind/SSD-1B) is a distilled, 50% smaller version of SDXL, offering up to 60% speedup
17
- '''
18
- if not torch.cuda.is_available():
19
- DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
20
 
21
  MAX_SEED = np.iinfo(np.int32).max
22
- CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
23
- MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
24
- USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "1") == "1"
25
- ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
26
- ENABLE_REFINER = os.getenv("ENABLE_REFINER", "0") == "1"
27
-
28
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
29
-
30
- style_list = [
31
- {
32
- "name": "(No style)",
33
- "prompt": "{prompt}",
34
- "negative_prompt": "",
35
- },
36
- {
37
- "name": "Cinematic",
38
- "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
39
- "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
40
- },
41
- {
42
- "name": "Photographic",
43
- "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed",
44
- "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
45
- },
46
- {
47
- "name": "Anime",
48
- "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed",
49
- "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
50
- },
51
- {
52
- "name": "Manga",
53
- "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style",
54
- "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic, Western comic style",
55
- },
56
- {
57
- "name": "Digital Art",
58
- "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
59
- "negative_prompt": "photo, photorealistic, realism, ugly",
60
- },
61
- {
62
- "name": "Pixel art",
63
- "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
64
- "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
65
- },
66
- {
67
- "name": "Fantasy art",
68
- "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
69
- "negative_prompt": "photographic, realistic, realism, 35mm film, dslr, cropped, frame, text, deformed, glitch, noise, noisy, off-center, deformed, cross-eyed, closed eyes, bad anatomy, ugly, disfigured, sloppy, duplicate, mutated, black and white",
70
- },
71
- {
72
- "name": "Neonpunk",
73
- "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
74
- "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured",
75
- },
76
- {
77
- "name": "3D Model",
78
- "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
79
- "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
80
- },
81
- ]
82
 
83
- styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
84
- STYLE_NAMES = list(styles.keys())
85
- DEFAULT_STYLE_NAME = "Cinematic"
86
 
 
 
 
 
 
 
 
 
 
87
 
88
- def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
89
- p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
90
- if not negative:
91
- negative = ""
92
- return p.replace("{prompt}", positive), n + negative
93
-
94
-
95
- if torch.cuda.is_available():
96
- vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
97
- pipe = StableDiffusionXLPipeline.from_pretrained(
98
- "segmind/SSD-1B",
99
- vae=vae,
100
- torch_dtype=torch.float16,
101
- use_safetensors=True,
102
- variant="fp16",
103
- )
104
- if ENABLE_REFINER:
105
- refiner = DiffusionPipeline.from_pretrained(
106
- "stabilityai/stable-diffusion-xl-refiner-1.0",
107
- vae=vae,
108
- torch_dtype=torch.float16,
109
- use_safetensors=True,
110
- variant="fp16",
111
- )
112
-
113
- if ENABLE_CPU_OFFLOAD:
114
- pipe.enable_model_cpu_offload()
115
- if ENABLE_REFINER:
116
- refiner.enable_model_cpu_offload()
117
- else:
118
- pipe.to(device)
119
- if ENABLE_REFINER:
120
- refiner.to(device)
121
- print("Loaded on Device!")
122
-
123
- if USE_TORCH_COMPILE:
124
- pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
125
- if ENABLE_REFINER:
126
- refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
127
- print("Model Compiled!")
128
 
129
  def save_image(img):
130
- unique_name = str(uuid.uuid4()) + '.png'
131
  img.save(unique_name)
132
  return unique_name
133
 
 
134
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
135
  if randomize_seed:
136
  seed = random.randint(0, MAX_SEED)
137
  return seed
138
 
 
 
139
  def generate(
140
  prompt: str,
141
- negative_prompt: str = "",
142
- style: str = DEFAULT_STYLE_NAME,
143
- prompt_2: str = "",
144
- negative_prompt_2: str = "",
145
- use_negative_prompt: bool = False,
146
- use_prompt_2: bool = False,
147
- use_negative_prompt_2: bool = False,
148
  seed: int = 0,
149
- width: int = 1024,
150
- height: int = 1024,
151
- guidance_scale_base: float = 5.0,
152
- guidance_scale_refiner: float = 5.0,
153
- num_inference_steps_base: int = 25,
154
- num_inference_steps_refiner: int = 25,
155
- apply_refiner: bool = False,
156
  randomize_seed: bool = False,
157
- progress = gr.Progress(track_tqdm=True)
158
  ):
159
  seed = randomize_seed_fn(seed, randomize_seed)
160
  generator = torch.Generator().manual_seed(seed)
161
 
162
- if not use_negative_prompt:
163
- negative_prompt = None # type: ignore
164
- if not use_prompt_2:
165
- prompt_2 = None # type: ignore
166
- if not use_negative_prompt_2:
167
- negative_prompt_2 = None # type: ignore
168
- prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
169
- if not apply_refiner:
170
- image = pipe(
171
- prompt=prompt,
172
- negative_prompt=negative_prompt,
173
- prompt_2=prompt_2,
174
- negative_prompt_2=negative_prompt_2,
175
- width=width,
176
- height=height,
177
- guidance_scale=guidance_scale_base,
178
- num_inference_steps=num_inference_steps_base,
179
- generator=generator,
180
- output_type="pil",
181
- ).images[0]
182
- else:
183
- latents = pipe(
184
- prompt=prompt,
185
- negative_prompt=negative_prompt,
186
- prompt_2=prompt_2,
187
- negative_prompt_2=negative_prompt_2,
188
- width=width,
189
- height=height,
190
- guidance_scale=guidance_scale_base,
191
- num_inference_steps=num_inference_steps_base,
192
- generator=generator,
193
- output_type="latent",
194
- ).images
195
- image = refiner(
196
- prompt=prompt,
197
- negative_prompt=negative_prompt,
198
- prompt_2=prompt_2,
199
- negative_prompt_2=negative_prompt_2,
200
- guidance_scale=guidance_scale_refiner,
201
- num_inference_steps=num_inference_steps_refiner,
202
- image=latents,
203
- generator=generator,
204
- ).images[0]
205
-
206
  image_path = save_image(image)
207
  print(image_path)
208
  return [image_path], seed
209
 
210
 
211
- examples = ['3d digital art of an adorable ghost, glowing within, holding a heart shaped pumpkin, Halloween, super cute, spooky haunted house background', 'beautiful lady, freckles, big smile, blue eyes, short ginger hair, dark makeup, wearing a floral blue vest top, soft light, dark grey background', 'professional portrait photo of an anthropomorphic cat wearing fancy gentleman hat and jacket walking in autumn forest.', 'an astronaut sitting in a diner, eating fries, cinematic, analog film', 'Albert Einstein in a surrealist Cyberpunk 2077 world, hyperrealistic', 'cinematic film still of Futuristic hero with golden dark armour with machine gun, muscular body']
 
 
 
 
 
212
 
213
  with gr.Blocks(css="style.css") as demo:
214
  gr.Markdown(DESCRIPTION)
@@ -229,34 +95,6 @@ with gr.Blocks(css="style.css") as demo:
229
  run_button = gr.Button("Run", scale=0)
230
  result = gr.Gallery(label="Result", columns=1, show_label=False)
231
  with gr.Accordion("Advanced options", open=False):
232
- with gr.Row():
233
- use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
234
- use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
235
- use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
236
- style_selection = gr.Radio(
237
- show_label=True, container=True, interactive=True,
238
- choices=STYLE_NAMES,
239
- value=DEFAULT_STYLE_NAME,
240
- label='Image Style'
241
- )
242
- negative_prompt = gr.Text(
243
- label="Negative prompt",
244
- max_lines=1,
245
- placeholder="Enter a negative prompt",
246
- visible=False,
247
- )
248
- prompt_2 = gr.Text(
249
- label="Prompt 2",
250
- max_lines=1,
251
- placeholder="Enter your prompt",
252
- visible=False,
253
- )
254
- negative_prompt_2 = gr.Text(
255
- label="Negative prompt 2",
256
- max_lines=1,
257
- placeholder="Enter a negative prompt",
258
- visible=False,
259
- )
260
  seed = gr.Slider(
261
  label="Seed",
262
  minimum=0,
@@ -280,36 +118,20 @@ with gr.Blocks(css="style.css") as demo:
280
  step=32,
281
  value=1024,
282
  )
283
- apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
284
  with gr.Row():
285
- guidance_scale_base = gr.Slider(
286
- label="Guidance scale for base",
287
- minimum=1,
288
- maximum=20,
289
- step=0.1,
290
- value=9.0,
291
- )
292
- num_inference_steps_base = gr.Slider(
293
- label="Number of inference steps for base",
294
- minimum=10,
295
- maximum=100,
296
- step=1,
297
- value=25,
298
- )
299
- with gr.Row(visible=False) as refiner_params:
300
- guidance_scale_refiner = gr.Slider(
301
- label="Guidance scale for refiner",
302
  minimum=1,
303
  maximum=20,
304
  step=0.1,
305
- value=5.0,
306
  )
307
- num_inference_steps_refiner = gr.Slider(
308
- label="Number of inference steps for refiner",
309
  minimum=10,
310
  maximum=100,
311
  step=1,
312
- value=25,
313
  )
314
 
315
  gr.Examples(
@@ -320,66 +142,16 @@ with gr.Blocks(css="style.css") as demo:
320
  cache_examples=CACHE_EXAMPLES,
321
  )
322
 
323
- use_negative_prompt.change(
324
- fn=lambda x: gr.update(visible=x),
325
- inputs=use_negative_prompt,
326
- outputs=negative_prompt,
327
- queue=False,
328
- api_name=False,
329
- )
330
- use_prompt_2.change(
331
- fn=lambda x: gr.update(visible=x),
332
- inputs=use_prompt_2,
333
- outputs=prompt_2,
334
- queue=False,
335
- api_name=False,
336
- )
337
- use_negative_prompt_2.change(
338
- fn=lambda x: gr.update(visible=x),
339
- inputs=use_negative_prompt_2,
340
- outputs=negative_prompt_2,
341
- queue=False,
342
- api_name=False,
343
- )
344
- apply_refiner.change(
345
- fn=lambda x: gr.update(visible=x),
346
- inputs=apply_refiner,
347
- outputs=refiner_params,
348
- queue=False,
349
- api_name=False,
350
- )
351
-
352
  gr.on(
353
  triggers=[
354
  prompt.submit,
355
- negative_prompt.submit,
356
- prompt_2.submit,
357
- negative_prompt_2.submit,
358
  run_button.click,
359
  ],
360
  fn=generate,
361
- inputs=[
362
- prompt,
363
- negative_prompt,
364
- style_selection,
365
- prompt_2,
366
- negative_prompt_2,
367
- use_negative_prompt,
368
- use_prompt_2,
369
- use_negative_prompt_2,
370
- seed,
371
- width,
372
- height,
373
- guidance_scale_base,
374
- guidance_scale_refiner,
375
- num_inference_steps_base,
376
- num_inference_steps_refiner,
377
- apply_refiner,
378
- randomize_seed
379
- ],
380
  outputs=[result, seed],
381
  api_name="run",
382
  )
383
 
384
  if __name__ == "__main__":
385
- demo.queue(max_size=20).launch()
 
 
 
 
 
1
  import os
2
  import random
3
 
4
  import gradio as gr
5
  import numpy as np
 
6
  import torch
7
+ from diffusers import DiffusionPipeline, UNet2DConditionModel
8
+ import spaces
9
  import uuid
10
 
11
+ DESCRIPTION = """# SPRIGHT T2I
12
+ #### [SPRIGHT T2I](https://spright.github.io/) is a framework to improve the spatial consistency of text-to-image models WITHOUT compromising their fidelity aspects.
13
+ """
 
 
14
 
15
  MAX_SEED = np.iinfo(np.int32).max
16
+ CACHE_EXAMPLES = os.getenv("CACHE_EXAMPLES", "1") == "1"
17
+ MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "768"))
18
+ TOKEN = os.getenv("HF_TOKEN")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
 
 
 
20
 
21
+ pipe_id = "SPRIGHT-T2I/spright-t2i-v1"
22
+ unet = UNet2DConditionModel.from_pretrained(pipe_id, subfolder="unet_ema", torch_dtype=torch.float16)
23
+ pipe = DiffusionPipeline.from_pretrained(
24
+ pipe_id,
25
+ unet=unet,
26
+ torch_dtype=torch.float16,
27
+ use_safetensors=True,
28
+ token=TOKEN,
29
+ ).to("cuda")
30
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  def save_image(img):
33
+ unique_name = str(uuid.uuid4()) + ".png"
34
  img.save(unique_name)
35
  return unique_name
36
 
37
+
38
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
39
  if randomize_seed:
40
  seed = random.randint(0, MAX_SEED)
41
  return seed
42
 
43
+
44
+ @spaces.gpu
45
  def generate(
46
  prompt: str,
 
 
 
 
 
 
 
47
  seed: int = 0,
48
+ width: int = 768,
49
+ height: int = 768,
50
+ guidance_scale: float = 7.5,
51
+ num_inference_steps: int = 50,
 
 
 
52
  randomize_seed: bool = False,
53
+ progress=gr.Progress(track_tqdm=True),
54
  ):
55
  seed = randomize_seed_fn(seed, randomize_seed)
56
  generator = torch.Generator().manual_seed(seed)
57
 
58
+ image = pipe(
59
+ prompt=prompt,
60
+ width=width,
61
+ height=height,
62
+ guidance_scale=guidance_scale,
63
+ num_inference_steps=num_inference_steps,
64
+ generator=generator,
65
+ ).images[0]
66
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
67
  image_path = save_image(image)
68
  print(image_path)
69
  return [image_path], seed
70
 
71
 
72
+ examples = [
73
+ "A cat next to a suitcase",
74
+ "A candle on the left of a mouse",
75
+ "A bag on the right of a dog",
76
+ "A mouse on the top of a bowl",
77
+ ]
78
 
79
  with gr.Blocks(css="style.css") as demo:
80
  gr.Markdown(DESCRIPTION)
 
95
  run_button = gr.Button("Run", scale=0)
96
  result = gr.Gallery(label="Result", columns=1, show_label=False)
97
  with gr.Accordion("Advanced options", open=False):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
  seed = gr.Slider(
99
  label="Seed",
100
  minimum=0,
 
118
  step=32,
119
  value=1024,
120
  )
 
121
  with gr.Row():
122
+ guidance_scale = gr.Slider(
123
+ label="Guidance scale",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  minimum=1,
125
  maximum=20,
126
  step=0.1,
127
+ value=7.5,
128
  )
129
+ num_inference_steps = gr.Slider(
130
+ label="Number of inference steps",
131
  minimum=10,
132
  maximum=100,
133
  step=1,
134
+ value=50,
135
  )
136
 
137
  gr.Examples(
 
142
  cache_examples=CACHE_EXAMPLES,
143
  )
144
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
145
  gr.on(
146
  triggers=[
147
  prompt.submit,
 
 
 
148
  run_button.click,
149
  ],
150
  fn=generate,
151
+ inputs=[prompt, seed, width, height, guidance_scale, num_inference_steps, randomize_seed],
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
152
  outputs=[result, seed],
153
  api_name="run",
154
  )
155
 
156
  if __name__ == "__main__":
157
+ demo.queue(max_size=20).launch()
requirements.txt CHANGED
@@ -1,7 +1,4 @@
1
- git+https://github.com/huggingface/diffusers.git
2
- accelerate==0.23.0
3
- gradio==3.45.2
4
- invisible-watermark==0.2.0
5
- Pillow==10.0.1
6
- torch==2.0.1
7
- transformers==4.33.3
 
1
+ diffusers
2
+ accelerate
3
+ torch
4
+ transformers