KingNish commited on
Commit
641e3d5
1 Parent(s): 6ea4158

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +96 -74
app.py CHANGED
@@ -1,76 +1,100 @@
1
  from __future__ import annotations
2
-
3
  import os
4
  import random
5
  import uuid
6
- from typing import Tuple
7
-
8
  import gradio as gr
 
9
  import numpy as np
 
 
10
  import torch
11
- from diffusers import LCMScheduler, PixArtAlphaPipeline
 
12
 
13
- # Use a more descriptive variable name
14
- MODEL_NAME = "PixArt-alpha/PixArt-LCM-XL-2-1024-MS"
15
 
16
- # Move environment variable checks and definitions to the top for better readability
17
- DESCRIPTION = """# Instant Image
18
  ### Super fast text to Image Generator.
19
  ### <span style='color: red;'>You may change the steps from 4 to 8, if you didn't get satisfied results.
20
  ### First Image processing takes time then images generate faster.
21
  """
 
 
22
 
 
 
23
  MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4192"))
24
  USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
25
  ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
26
  PORT = int(os.getenv("DEMO_PORT", "15432"))
27
 
28
- # Check CUDA availability early on
29
- if not torch.cuda.is_available():
30
- DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
31
-
32
- # Cache examples only if CUDA is available
33
- CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
34
-
35
- MAX_SEED = np.iinfo(np.int32).max
36
- NUM_IMAGES_PER_PROMPT = 1
37
 
38
- # Use Enum for better style management
39
- from enum import Enum
40
 
41
- class Style(Enum):
42
- NO_STYLE = ("(No style)", "{prompt}", "")
43
- CINEMATIC = ("Cinematic", "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy", "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured")
44
- REALISTIC = ("Realistic", "Photorealistic {prompt} . Ulta-realistic, professional, 4k, highly detailed", "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, disfigured")
45
- ANIME = ("Anime", "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed", "photo, deformed, black and white, realism, disfigured, low contrast")
46
- DIGITAL_ART = ("Digital Art", "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed", "photo, photorealistic, realism, ugly")
47
- PIXEL_ART = ("Pixel art", "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics", "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic")
48
- FANTASY_ART = ("Fantasy art", "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy", "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")
49
- THREE_D_MODEL = ("3D Model", "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting", "ugly, deformed, noisy, low poly, blurry, painting")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
 
51
- def __init__(self, name, prompt, negative_prompt):
52
- self.name = name
53
- self.prompt = prompt
54
- self.negative_prompt = negative_prompt
55
 
56
- # Use the Enum values directly
57
- styles = {style.name: (style.prompt, style.negative_prompt) for style in Style}
58
  STYLE_NAMES = list(styles.keys())
59
- DEFAULT_STYLE_NAME = Style.NO_STYLE.name
 
60
 
61
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
 
 
62
 
63
- # Load the pipeline only if CUDA is available
64
  if torch.cuda.is_available():
 
65
  pipe = PixArtAlphaPipeline.from_pretrained(
66
- MODEL_NAME,
67
  torch_dtype=torch.float16,
68
  use_safetensors=True,
69
  )
70
 
71
  if os.getenv('CONSISTENCY_DECODER', False):
72
  print("Using DALL-E 3 Consistency Decoder")
73
- # Assuming ConsistencyDecoderVAE is defined elsewhere
74
  pipe.vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
75
 
76
  if ENABLE_CPU_OFFLOAD:
@@ -78,8 +102,8 @@ if torch.cuda.is_available():
78
  else:
79
  pipe.to(device)
80
  print("Loaded on Device!")
81
-
82
- # Speed-up T5
83
  pipe.text_encoder.to_bettertransformer()
84
 
85
  if USE_TORCH_COMPILE:
@@ -87,40 +111,37 @@ if torch.cuda.is_available():
87
  print("Model Compiled!")
88
 
89
  def save_image(img):
90
- # Generate image names in a temporary directory
91
- os.makedirs("tmp", exist_ok=True)
92
- unique_name = os.path.join("tmp", f"{uuid.uuid4()}.png")
93
  img.save(unique_name)
94
  return unique_name
95
 
 
 
96
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
97
- return random.randint(0, MAX_SEED) if randomize_seed else seed
 
 
98
 
99
- # No need to use @spaces.GPU if you're checking CUDA availability within the function
100
  def generate(
101
- prompt: str,
102
- negative_prompt: str = "",
103
- style: str = DEFAULT_STYLE_NAME,
104
- use_negative_prompt: bool = False,
105
- seed: int = 0,
106
- width: int = 1024,
107
- height: int = 1024,
108
- inference_steps: int = 8,
109
- randomize_seed: bool = False,
110
- use_resolution_binning: bool = True,
111
- progress=gr.Progress(track_tqdm=True),
112
  ):
113
- if not torch.cuda.is_available():
114
- return "This demo requires a GPU to run.", seed
115
-
116
  seed = int(randomize_seed_fn(seed, randomize_seed))
117
  generator = torch.Generator().manual_seed(seed)
118
-
119
  if not use_negative_prompt:
120
- negative_prompt = None
121
-
122
- prompt, negative_prompt = styles.get(style, styles[DEFAULT_STYLE_NAME])
123
- prompt = prompt.replace("{prompt}", prompt)
124
 
125
  images = pipe(
126
  prompt=prompt,
@@ -136,6 +157,7 @@ def generate(
136
  ).images
137
 
138
  image_paths = [save_image(img) for img in images]
 
139
  return image_paths, seed
140
 
141
 
@@ -170,17 +192,17 @@ with gr.Blocks(css=css) as demo:
170
  )
171
  run_button = gr.Button("Run", scale=0)
172
  result = gr.Gallery(label="Result", columns=1, show_label=False)
173
-
174
  with gr.Accordion("Advanced options", open=False):
175
  with gr.Group():
176
  with gr.Row():
177
  use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False, visible=True)
178
  negative_prompt = gr.Text(
179
- label="Negative prompt",
180
- max_lines=1,
181
- placeholder="Enter a negative prompt",
182
- visible=True,
183
- )
184
  style_selection = gr.Radio(
185
  show_label=True,
186
  container=True,
@@ -195,7 +217,7 @@ with gr.Blocks(css=css) as demo:
195
  maximum=MAX_SEED,
196
  step=1,
197
  value=0,
198
- )
199
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
200
  with gr.Row(visible=True):
201
  width = gr.Slider(
@@ -242,7 +264,7 @@ with gr.Blocks(css=css) as demo:
242
  run_button.click,
243
  ],
244
  batch=True,
245
- max_batch_size=10,
246
  fn=generate,
247
  inputs=[
248
  prompt,
@@ -260,4 +282,4 @@ with gr.Blocks(css=css) as demo:
260
  )
261
 
262
  if __name__ == "__main__":
263
- demo.queue(max_size=200).launch(server_port=PORT)
 
1
  from __future__ import annotations
 
2
  import os
3
  import random
4
  import uuid
 
 
5
  import gradio as gr
6
+ import spaces
7
  import numpy as np
8
+ import uuid
9
+ from diffusers import PixArtAlphaPipeline, LCMScheduler
10
  import torch
11
+ from typing import Tuple
12
+ from datetime import datetime
13
 
 
 
14
 
15
+ DESCRIPTION = """ # Instant Image
 
16
  ### Super fast text to Image Generator.
17
  ### <span style='color: red;'>You may change the steps from 4 to 8, if you didn't get satisfied results.
18
  ### First Image processing takes time then images generate faster.
19
  """
20
+ if not torch.cuda.is_available():
21
+ DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
22
 
23
+ MAX_SEED = np.iinfo(np.int32).max
24
+ CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"
25
  MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4192"))
26
  USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
27
  ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
28
  PORT = int(os.getenv("DEMO_PORT", "15432"))
29
 
30
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 
 
 
 
 
 
 
31
 
 
 
32
 
33
+ style_list = [
34
+ {
35
+ "name": "(No style)",
36
+ "prompt": "{prompt}",
37
+ "negative_prompt": "",
38
+ },
39
+ {
40
+ "name": "Cinematic",
41
+ "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
42
+ "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
43
+ },
44
+ {
45
+ "name": "Realistic",
46
+ "prompt": "Photorealistic {prompt} . Ulta-realistic, professional, 4k, highly detailed",
47
+ "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, disfigured",
48
+ },
49
+ {
50
+ "name": "Anime",
51
+ "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime, highly detailed",
52
+ "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
53
+ },
54
+ {
55
+ "name": "Digital Art",
56
+ "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed",
57
+ "negative_prompt": "photo, photorealistic, realism, ugly",
58
+ },
59
+ {
60
+ "name": "Pixel art",
61
+ "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics",
62
+ "negative_prompt": "sloppy, messy, blurry, noisy, highly detailed, ultra textured, photo, realistic",
63
+ },
64
+ {
65
+ "name": "Fantasy art",
66
+ "prompt": "ethereal fantasy concept art of {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy",
67
+ "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",
68
+ },
69
+ {
70
+ "name": "3D Model",
71
+ "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
72
+ "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
73
+ },
74
+ ]
75
 
 
 
 
 
76
 
77
+ styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 
78
  STYLE_NAMES = list(styles.keys())
79
+ DEFAULT_STYLE_NAME = "(No style)"
80
+ NUM_IMAGES_PER_PROMPT = 1
81
 
82
+ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
83
+ p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
84
+ if not negative:
85
+ negative = ""
86
+ return p.replace("{prompt}", positive), n + negative
87
 
 
88
  if torch.cuda.is_available():
89
+
90
  pipe = PixArtAlphaPipeline.from_pretrained(
91
+ "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
92
  torch_dtype=torch.float16,
93
  use_safetensors=True,
94
  )
95
 
96
  if os.getenv('CONSISTENCY_DECODER', False):
97
  print("Using DALL-E 3 Consistency Decoder")
 
98
  pipe.vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=torch.float16)
99
 
100
  if ENABLE_CPU_OFFLOAD:
 
102
  else:
103
  pipe.to(device)
104
  print("Loaded on Device!")
105
+
106
+ # speed-up T5
107
  pipe.text_encoder.to_bettertransformer()
108
 
109
  if USE_TORCH_COMPILE:
 
111
  print("Model Compiled!")
112
 
113
  def save_image(img):
114
+ unique_name = str(uuid.uuid4()) + ".png"
 
 
115
  img.save(unique_name)
116
  return unique_name
117
 
118
+
119
+
120
  def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
121
+ if randomize_seed:
122
+ seed = random.randint(0, MAX_SEED)
123
+ return seed
124
 
125
+ @spaces.GPU(duration=30)
126
  def generate(
127
+ prompt: str,
128
+ negative_prompt: str = "",
129
+ style: str = DEFAULT_STYLE_NAME,
130
+ use_negative_prompt: bool = False,
131
+ seed: int = 0,
132
+ width: int = 1024,
133
+ height: int = 1024,
134
+ inference_steps: int = 8,
135
+ randomize_seed: bool = False,
136
+ use_resolution_binning: bool = True,
137
+ progress=gr.Progress(track_tqdm=True),
138
  ):
 
 
 
139
  seed = int(randomize_seed_fn(seed, randomize_seed))
140
  generator = torch.Generator().manual_seed(seed)
141
+
142
  if not use_negative_prompt:
143
+ negative_prompt = None # type: ignore
144
+ prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
 
 
145
 
146
  images = pipe(
147
  prompt=prompt,
 
157
  ).images
158
 
159
  image_paths = [save_image(img) for img in images]
160
+ print(image_paths)
161
  return image_paths, seed
162
 
163
 
 
192
  )
193
  run_button = gr.Button("Run", scale=0)
194
  result = gr.Gallery(label="Result", columns=1, show_label=False)
195
+
196
  with gr.Accordion("Advanced options", open=False):
197
  with gr.Group():
198
  with gr.Row():
199
  use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False, visible=True)
200
  negative_prompt = gr.Text(
201
+ label="Negative prompt",
202
+ max_lines=1,
203
+ placeholder="Enter a negative prompt",
204
+ visible=True,
205
+ )
206
  style_selection = gr.Radio(
207
  show_label=True,
208
  container=True,
 
217
  maximum=MAX_SEED,
218
  step=1,
219
  value=0,
220
+ )
221
  randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
222
  with gr.Row(visible=True):
223
  width = gr.Slider(
 
264
  run_button.click,
265
  ],
266
  batch=True,
267
+ max_batch_size=10,
268
  fn=generate,
269
  inputs=[
270
  prompt,
 
282
  )
283
 
284
  if __name__ == "__main__":
285
+ demo.queue(max_size=200).launch()