prithivMLmods commited on
Commit
ee3de9b
·
verified ·
1 Parent(s): c3a087b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +287 -0
app.py ADDED
@@ -0,0 +1,287 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import gradio as gr
3
+ import numpy as np
4
+ from PIL import Image
5
+ import spaces
6
+ import torch
7
+ from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
8
+ import os
9
+ import uuid
10
+ import random
11
+
12
+ # Description for the Gradio interface
13
+ DESCRIPTIONx = """## INSTANT WALLPAPER """
14
+
15
+ # CSS for styling the Gradio interface
16
+ css = '''
17
+ .gradio-container{max-width: 575px !important}
18
+ h1{text-align:center}
19
+ footer {
20
+ visibility: hidden
21
+ }
22
+ '''
23
+
24
+ # Example prompts for the user to try
25
+ examples = [
26
+ "Illustration of A starry night camp in the mountains. Low-angle view, Minimal background, Geometric shapes theme, Pottery, Split-complementary colors, Bicolored light, UHD",
27
+ "Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic oil --ar 2:3 --q 2 --s 750 --v 5 --ar 2:3 --q 2 --s 750 --v 5"
28
+ ]
29
+
30
+ # Environment variables and defaults for configuration
31
+ MODEL_ID = os.getenv("MODEL_USED") #SG161222/RealVisXL_V4.0 / SG161222/Realistic_Vision_V5.1_noVAE / SG161222/RealVisXL_V4.0_Lightning (1/3)
32
+ MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
33
+ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
34
+ ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
35
+ BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))
36
+
37
+ # Setting the device to GPU if available, otherwise CPU
38
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
39
+
40
+ # Loading the Stable Diffusion model
41
+ pipe = StableDiffusionXLPipeline.from_pretrained(
42
+ MODEL_ID,
43
+ torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
44
+ use_safetensors=True,
45
+ add_watermarker=False,
46
+ ).to(device)
47
+
48
+ # Configuring the scheduler for the model
49
+ pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
50
+
51
+ # Compiling the model for performance improvement if enabled
52
+ if USE_TORCH_COMPILE:
53
+ pipe.compile()
54
+
55
+ # Enabling CPU offload to save GPU memory if enabled
56
+ if ENABLE_CPU_OFFLOAD:
57
+ pipe.enable_model_cpu_offload()
58
+
59
+ # Maximum seed value for randomization
60
+ MAX_SEED = np.iinfo(np.int32).max
61
+
62
+ # Function to save the generated image
63
+ def save_image(img):
64
+ unique_name = str(uuid.uuid4()) + ".png"
65
+ img.save(unique_name)
66
+ return unique_name
67
+
68
+ # Function to randomize the seed if needed
69
+ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
70
+ if randomize_seed:
71
+ seed = random.randint(0, MAX_SEED)
72
+ return seed
73
+
74
+ # Defining the main generation function with GPU acceleration
75
+ @spaces.GPU(duration=60, enable_queue=True)
76
+ def generate(
77
+ prompt: str,
78
+ negative_prompt: str = "",
79
+ use_negative_prompt: bool = False,
80
+ seed: int = 1,
81
+ width: int = 1024,
82
+ height: int = 1024,
83
+ guidance_scale: float = 3,
84
+ num_inference_steps: int = 25,
85
+ randomize_seed: bool = False,
86
+ use_resolution_binning: bool = True,
87
+ num_images: int = 1,
88
+ progress=gr.Progress(track_tqdm=True),
89
+ ):
90
+ # Randomizing the seed if required
91
+ seed = int(randomize_seed_fn(seed, randomize_seed))
92
+ generator = torch.Generator(device=device).manual_seed(seed)
93
+
94
+ # Setting up the options for the image generation
95
+ options = {
96
+ "prompt": [prompt] * num_images,
97
+ "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
98
+ "width": width,
99
+ "height": height,
100
+ "guidance_scale": guidance_scale,
101
+ "num_inference_steps": num_inference_steps,
102
+ "generator": generator,
103
+ "output_type": "pil",
104
+ }
105
+
106
+ if use_resolution_binning:
107
+ options["use_resolution_binning"] = True
108
+
109
+ # Generating images in batches
110
+ images = []
111
+ for i in range(0, num_images, BATCH_SIZE):
112
+ batch_options = options.copy()
113
+ batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
114
+ if "negative_prompt" in batch_options:
115
+ batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
116
+ images.extend(pipe(**batch_options).images)
117
+
118
+ # Saving the generated images
119
+ image_paths = [save_image(img) for img in images]
120
+ return image_paths, seed
121
+
122
+ # Function to set the wallpaper size based on the selected option
123
+ def set_wallpaper_size(size):
124
+ if size == "phone":
125
+ return 1080, 1920
126
+ elif size == "desktop":
127
+ return 1920, 1080
128
+ return 1024, 1024
129
+
130
+ # Function to load predefined images for display
131
+ def load_predefined_images():
132
+ predefined_images = [
133
+ "assets/image1.png",
134
+ "assets/image2.png",
135
+ "assets/image3.png",
136
+ "assets/image4.png",
137
+ "assets/image5.png",
138
+ "assets/image6.png",
139
+ "assets/image7.png",
140
+ "assets/image8.png",
141
+ "assets/image9.png",
142
+ ]
143
+ return predefined_images
144
+
145
+ # Defining the Gradio interface with blocks
146
+ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
147
+ gr.Markdown(DESCRIPTIONx)
148
+ with gr.Group():
149
+ with gr.Row():
150
+ prompt = gr.Text(
151
+ label="Prompt",
152
+ show_label=False,
153
+ max_lines=1,
154
+ placeholder="Enter your prompt",
155
+ container=False,
156
+ )
157
+ run_button = gr.Button("Run", scale=0)
158
+ result = gr.Gallery(label="Result", columns=1, show_label=False)
159
+
160
+ with gr.Group():
161
+ wallpaper_size = gr.Radio(
162
+ choices=["phone", "desktop", "custom"],
163
+ label="Wallpaper Size",
164
+ value="desktop"
165
+ )
166
+ width = gr.Slider(
167
+ label="Width",
168
+ minimum=512,
169
+ maximum=MAX_IMAGE_SIZE,
170
+ step=64,
171
+ value=1920,
172
+ visible=False,
173
+ )
174
+ height = gr.Slider(
175
+ label="Height",
176
+ minimum=512,
177
+ maximum=MAX_IMAGE_SIZE,
178
+ step=64,
179
+ value=1080,
180
+ visible=False,
181
+ )
182
+
183
+ # Changing the wallpaper size based on user selection
184
+ wallpaper_size.change(
185
+ fn=set_wallpaper_size,
186
+ inputs=wallpaper_size,
187
+ outputs=[width, height],
188
+ api_name="set_wallpaper_size"
189
+ )
190
+
191
+ # Advanced options for image generation
192
+ with gr.Accordion("Advanced options", open=False, visible=False):
193
+ num_images = gr.Slider(
194
+ label="Number of Images",
195
+ minimum=1,
196
+ maximum=4,
197
+ step=1,
198
+ value=1,
199
+ )
200
+ with gr.Row():
201
+ use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
202
+ negative_prompt = gr.Text(
203
+ label="Negative prompt",
204
+ max_lines=5,
205
+ lines=4,
206
+ placeholder="Enter a negative prompt",
207
+ 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",
208
+ visible=True,
209
+ )
210
+ seed = gr.Slider(
211
+ label="Seed",
212
+ minimum=0,
213
+ maximum=MAX_SEED,
214
+ step=1,
215
+ value=0,
216
+ )
217
+ randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
218
+ with gr.Row():
219
+ guidance_scale = gr.Slider(
220
+ label="Guidance Scale",
221
+ minimum=0.1,
222
+ maximum=6,
223
+ step=0.1,
224
+ value=3.0,
225
+ )
226
+ num_inference_steps = gr.Slider(
227
+ label="Number of inference steps",
228
+ minimum=1,
229
+ maximum=25,
230
+ step=1,
231
+ value=20,
232
+ )
233
+
234
+ # Adding examples for the user to try
235
+ gr.Examples(
236
+ examples=examples,
237
+ inputs=prompt,
238
+ cache_examples=False
239
+ )
240
+
241
+ # Changing the visibility of the negative prompt based on user selection
242
+ use_negative_prompt.change(
243
+ fn=lambda x: gr.update(visible=x),
244
+ inputs=use_negative_prompt,
245
+ outputs=negative_prompt,
246
+ api_name=False,
247
+ )
248
+
249
+ # Setting up the triggers and linking them to the generate function
250
+ gr.on(
251
+ triggers=[
252
+ prompt.submit,
253
+ negative_prompt.submit,
254
+ run_button.click,
255
+ ],
256
+ fn=generate,
257
+ inputs=[
258
+ prompt,
259
+ negative_prompt,
260
+ use_negative_prompt,
261
+ seed,
262
+ width,
263
+ height,
264
+ guidance_scale,
265
+ num_inference_steps,
266
+ randomize_seed,
267
+ num_images
268
+ ],
269
+ outputs=[result, seed],
270
+ api_name="run",
271
+ )
272
+
273
+ # Adding a predefined gallery section
274
+ gr.Markdown("### Sample Images")
275
+ predefined_gallery = gr.Gallery(label="Predefined Images", columns=3, show_label=False, value=load_predefined_images())
276
+
277
+ # Adding a disclaimer
278
+ gr.Markdown("**Disclaimer:**")
279
+ gr.Markdown("This is the demo space for generating wallpapers using detailed prompts. This space works best for desktop-sized images (1920x1080). Reasonable quality images can be generated for mobile sizes (1080x1920), and custom images (1024x1024) can also be generated with better quality. Mobile settings may become disfigured. Try the sample prompts for generating higher quality images.<a href='https://huggingface.co/spaces/prithivMLmods/INSTANT-WALLPAPER/blob/main/sample_prompts.txt' target='_blank'>Try prompts</a>.")
280
+
281
+ # Adding a note about user responsibility
282
+ gr.Markdown("**Note:**")
283
+ gr.Markdown("⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.")
284
+
285
+ # Launching the Gradio interface
286
+ if __name__ == "__main__":
287
+ demo.queue(max_size=40).launch()