prithivMLmods commited on
Commit
4c1a95a
·
verified ·
1 Parent(s): bd9d3b6

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +409 -0
app.py ADDED
@@ -0,0 +1,409 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #path1.0398
2
+ import os
3
+ import random
4
+ import uuid
5
+ import json
6
+
7
+ import gradio as gr
8
+ import numpy as np
9
+ from PIL import Image
10
+ import spaces
11
+ import torch
12
+ from diffusers import DiffusionPipeline
13
+ from typing import Tuple
14
+
15
+ # BaseConditions
16
+ bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
17
+ bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
18
+ default_negative = os.getenv("default_negative","")
19
+
20
+ def check_text(prompt, negative=""):
21
+ for i in bad_words:
22
+ if i in prompt:
23
+ return True
24
+ for i in bad_words_negative:
25
+ if i in negative:
26
+ return True
27
+ return False
28
+
29
+ style_list = [
30
+ {
31
+ "name": "3840 x 2160",
32
+ "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
33
+ "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
34
+ },
35
+ {
36
+ "name": "2560 x 1440",
37
+ "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
38
+ "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
39
+ },
40
+ {
41
+ "name": "3D Model",
42
+ "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
43
+ "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
44
+ },
45
+ ]
46
+
47
+ collage_style_list = [
48
+
49
+
50
+ {
51
+ "name": "B & W",
52
+ "prompt": "black and white collage of {prompt}. monochromatic, timeless, classic, dramatic contrast",
53
+ "negative_prompt": "colorful, vibrant, bright, flashy",
54
+ },
55
+
56
+ {
57
+ "name": "Polaroid",
58
+ "prompt": "collage of polaroid photos featuring {prompt}. vintage style, high contrast, nostalgic, instant film aesthetic",
59
+ "negative_prompt": "digital, modern, low quality, blurry",
60
+ },
61
+
62
+ {
63
+ "name": "Watercolor",
64
+ "prompt": "watercolor collage of {prompt}. soft edges, translucent colors, painterly effects",
65
+ "negative_prompt": "digital, sharp lines, solid colors",
66
+ },
67
+
68
+ {
69
+ "name": "Cinematic",
70
+ "prompt": "cinematic collage of {prompt}. film stills, movie posters, dramatic lighting",
71
+ "negative_prompt": "static, lifeless, mundane",
72
+ },
73
+
74
+ {
75
+ "name": "Nostalgic",
76
+ "prompt": "nostalgic collage of {prompt}. retro imagery, vintage objects, sentimental journey",
77
+ "negative_prompt": "contemporary, futuristic, forward-looking",
78
+ },
79
+
80
+ {
81
+ "name": "Vintage",
82
+ "prompt": "vintage collage of {prompt}. aged paper, sepia tones, retro imagery, antique vibes",
83
+ "negative_prompt": "modern, contemporary, futuristic, high-tech",
84
+ },
85
+
86
+ {
87
+ "name": "Scrapbook",
88
+ "prompt": "scrapbook style collage of {prompt}. mixed media, hand-cut elements, textures, paper, stickers, doodles",
89
+ "negative_prompt": "clean, digital, modern, low quality",
90
+ },
91
+
92
+ {
93
+ "name": "NeoNGlow",
94
+ "prompt": "neon glow collage of {prompt}. vibrant colors, glowing effects, futuristic vibes",
95
+ "negative_prompt": "dull, muted colors, vintage, retro",
96
+ },
97
+
98
+ {
99
+ "name": "Geometric",
100
+ "prompt": "geometric collage of {prompt}. abstract shapes, colorful, sharp edges, modern design, high quality",
101
+ "negative_prompt": "blurry, low quality, traditional, dull",
102
+ },
103
+ {
104
+ "name": "Thematic",
105
+ "prompt": "thematic collage of {prompt}. cohesive theme, well-organized, matching colors, creative layout",
106
+ "negative_prompt": "random, messy, unorganized, clashing colors",
107
+ },
108
+
109
+ {
110
+ "name": "Retro Pop",
111
+ "prompt": "retro pop art collage of {prompt}. bold colors, comic book style, halftone dots, vintage ads",
112
+ "negative_prompt": "subdued colors, minimalist, modern, subtle",
113
+ },
114
+
115
+
116
+ {
117
+ "name": "No Style",
118
+ "prompt": "{prompt}",
119
+ "negative_prompt": "",
120
+ },
121
+ ]
122
+
123
+ styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
124
+ collage_styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in collage_style_list}
125
+ STYLE_NAMES = list(styles.keys())
126
+ COLLAGE_STYLE_NAMES = list(collage_styles.keys())
127
+ DEFAULT_STYLE_NAME = "3840 x 2160"
128
+ DEFAULT_COLLAGE_STYLE_NAME = "B & W"
129
+
130
+ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
131
+ if style_name in styles:
132
+ p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
133
+ elif style_name in collage_styles:
134
+ p, n = collage_styles.get(style_name, collage_styles[DEFAULT_COLLAGE_STYLE_NAME])
135
+ else:
136
+ p, n = styles[DEFAULT_STYLE_NAME]
137
+
138
+ if not negative:
139
+ negative = ""
140
+ return p.replace("{prompt}", positive), n + negative
141
+
142
+ DESCRIPTION = """"""
143
+ if not torch.cuda.is_available():
144
+ DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
145
+
146
+ MAX_SEED = np.iinfo(np.int32).max
147
+ CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
148
+ MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
149
+ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
150
+ ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
151
+
152
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
153
+
154
+ if torch.cuda.is_available():
155
+ pipe = DiffusionPipeline.from_pretrained(
156
+ "SG161222/RealVisXL_V4.0",
157
+ torch_dtype=torch.float16,
158
+ use_safetensors=True,
159
+ add_watermarker=False,
160
+ variant="fp16"
161
+ ).to(device)
162
+
163
+ if ENABLE_CPU_OFFLOAD:
164
+ pipe.enable_model_cpu_offload()
165
+ else:
166
+ pipe.to(device)
167
+ print("Loaded on Device!")
168
+
169
+ if USE_TORCH_COMPILE:
170
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
171
+ print("Model Compiled!")
172
+
173
+ def save_image(img, path):
174
+ img.save(path)
175
+
176
+ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
177
+ if randomize_seed:
178
+ seed = random.randint(0, MAX_SEED)
179
+ return seed
180
+
181
+ @spaces.GPU(enable_queue=True)
182
+ def generate(
183
+ prompt: str,
184
+ negative_prompt: str = "",
185
+ use_negative_prompt: bool = False,
186
+ style: str = DEFAULT_STYLE_NAME,
187
+ collage_style: str = DEFAULT_COLLAGE_STYLE_NAME,
188
+ grid_size: str = "2x2",
189
+ seed: int = 0,
190
+ width: int = 1024,
191
+ height: int = 1024,
192
+ guidance_scale: float = 3,
193
+ randomize_seed: bool = False,
194
+ use_resolution_binning: bool = True,
195
+ progress=gr.Progress(track_tqdm=True),
196
+ ):
197
+ if check_text(prompt, negative_prompt):
198
+ raise ValueError("Prompt contains restricted words.")
199
+
200
+ if collage_style != "No Style":
201
+ prompt, negative_prompt = apply_style(collage_style, prompt, negative_prompt)
202
+ else:
203
+ prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
204
+
205
+ seed = int(randomize_seed_fn(seed, randomize_seed))
206
+ generator = torch.Generator().manual_seed(seed)
207
+
208
+ if not use_negative_prompt:
209
+ negative_prompt = "" # type: ignore
210
+ negative_prompt += default_negative
211
+
212
+ grid_sizes = {
213
+ "2x1": (2, 1),
214
+ "1x2": (1, 2),
215
+ "2x2": (2, 2),
216
+ "2x3": (2, 3),
217
+ "3x2": (3, 2),
218
+ "1x1": (1, 1)
219
+ }
220
+
221
+ grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
222
+ num_images = grid_size_x * grid_size_y
223
+
224
+ options = {
225
+ "prompt": prompt,
226
+ "negative_prompt": negative_prompt,
227
+ "width": width,
228
+ "height": height,
229
+ "guidance_scale": guidance_scale,
230
+ "num_inference_steps": 20,
231
+ "generator": generator,
232
+ "num_images_per_prompt": num_images,
233
+ "use_resolution_binning": use_resolution_binning,
234
+ "output_type": "pil",
235
+ }
236
+
237
+ torch.cuda.empty_cache() # Clear GPU memory
238
+ images = pipe(**options).images
239
+
240
+ grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
241
+
242
+ for i, img in enumerate(images[:num_images]):
243
+ grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
244
+
245
+ unique_name = str(uuid.uuid4()) + ".png"
246
+ save_image(grid_img, unique_name)
247
+ return [unique_name], seed
248
+
249
+ examples = [
250
+
251
+ "Portrait of a beautiful woman in a hat, summer outfit, with freckles on her face, in a close up shot, with sunlight, outdoors, in soft light, with a beach background, looking at the camera, with high resolution photography, in the style of Hasselblad X2D50c --ar 85:128 --v 6.0 --style raw",
252
+ "Flying food photography with [Two Burgers] as the main subject, Splashes of Toppings and Seasonings, [Rocket Lettuce], [Cheddar Flavored Cheese], [Onion], [Pickles], [Special Sauce], [Sesame Bun], [ sea salt crystals] ::3 Capturing the dynamic splashes of food using high-speed photography , photorealistic, surrealism style, [white background], trending background [clean], Minimalist ::2 [Cuware], [Table], [ Steam], [Smoke], [Vegetable Leaves], [Tomato] ::-0.5 Ad Posters, Pro-Grade Color Grading, Studio Lighting, Rim Lights, [Layered Comps], EOS-1D X Mark III, 500px, Behance, concept art"
253
+
254
+ ]
255
+
256
+ css = '''
257
+ .gradio-container{max-width: 560px !important}
258
+ h1{text-align:center}
259
+ '''
260
+ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
261
+ gr.Markdown(DESCRIPTION)
262
+ gr.DuplicateButton(
263
+ value="Duplicate Space for private use",
264
+ elem_id="duplicate-button",
265
+ visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
266
+ )
267
+ with gr.Group():
268
+ with gr.Row():
269
+ prompt = gr.Text(
270
+ label="Prompt",
271
+ show_label=False,
272
+ max_lines=1,
273
+ placeholder="Enter your prompt",
274
+ container=False,
275
+ )
276
+ run_button = gr.Button("Run")
277
+ result = gr.Gallery(label="Grid", columns=1, preview=True)
278
+
279
+ with gr.Row(visible=True):
280
+ collage_style_selection = gr.Radio(
281
+ show_label=True,
282
+ container=True,
283
+ interactive=True,
284
+ choices=COLLAGE_STYLE_NAMES,
285
+ value=DEFAULT_COLLAGE_STYLE_NAME,
286
+ label="Collage Template",
287
+ )
288
+ with gr.Row(visible=True):
289
+ grid_size_selection = gr.Dropdown(
290
+ choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
291
+ value="2x2",
292
+ label="Grid Size"
293
+ )
294
+ with gr.Row(visible=True):
295
+ style_selection = gr.Radio(
296
+ show_label=True,
297
+ container=True,
298
+ interactive=True,
299
+ choices=STYLE_NAMES,
300
+ value=DEFAULT_STYLE_NAME,
301
+ label="Style",
302
+ )
303
+
304
+ with gr.Accordion("Advanced options", open=False):
305
+ use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
306
+ negative_prompt = gr.Text(
307
+ label="Negative prompt",
308
+ max_lines=1,
309
+ placeholder="Enter a negative prompt",
310
+ value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
311
+ visible=True,
312
+ )
313
+ with gr.Row():
314
+ num_inference_steps = gr.Slider(
315
+ label="Steps",
316
+ minimum=10,
317
+ maximum=30,
318
+ step=1,
319
+ value=15,
320
+ )
321
+ with gr.Row():
322
+ num_images_per_prompt = gr.Slider(
323
+ label="Images",
324
+ minimum=1,
325
+ maximum=5,
326
+ step=1,
327
+ value=2,
328
+ )
329
+ seed = gr.Slider(
330
+ label="Seed",
331
+ minimum=0,
332
+ maximum=MAX_SEED,
333
+ step=1,
334
+ value=0,
335
+ visible=True
336
+ )
337
+ randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
338
+
339
+
340
+ with gr.Row(visible=True):
341
+ width = gr.Slider(
342
+ label="Width",
343
+ minimum=512,
344
+ maximum=2048,
345
+ step=8,
346
+ value=1024,
347
+ )
348
+ height = gr.Slider(
349
+ label="Height",
350
+ minimum=512,
351
+ maximum=2048,
352
+ step=8,
353
+ value=1024,
354
+ )
355
+
356
+
357
+
358
+ with gr.Row():
359
+ guidance_scale = gr.Slider(
360
+ label="Guidance Scale",
361
+ minimum=0.1,
362
+ maximum=20.0,
363
+ step=0.1,
364
+ value=6,
365
+ )
366
+
367
+
368
+
369
+ gr.Examples(
370
+ examples=examples,
371
+ inputs=prompt,
372
+ outputs=[result, seed],
373
+ fn=generate,
374
+ cache_examples=CACHE_EXAMPLES,
375
+ )
376
+
377
+ use_negative_prompt.change(
378
+ fn=lambda x: gr.update(visible=x),
379
+ inputs=use_negative_prompt,
380
+ outputs=negative_prompt,
381
+ api_name=False,
382
+ )
383
+
384
+ gr.on(
385
+ triggers=[
386
+ prompt.submit,
387
+ negative_prompt.submit,
388
+ run_button.click,
389
+ ],
390
+ fn=generate,
391
+ inputs=[
392
+ prompt,
393
+ negative_prompt,
394
+ use_negative_prompt,
395
+ style_selection,
396
+ collage_style_selection,
397
+ grid_size_selection,
398
+ seed,
399
+ width,
400
+ height,
401
+ guidance_scale,
402
+ randomize_seed,
403
+ ],
404
+ outputs=[result, seed],
405
+ api_name="run",
406
+ )
407
+
408
+ if __name__ == "__main__":
409
+ demo.queue(max_size=20).launch()