ford442 commited on
Commit
9eddcb4
·
1 Parent(s): 98331f6

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +77 -2
app.py CHANGED
@@ -158,6 +158,57 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
158
  seed = random.randint(0, MAX_SEED)
159
  return seed
160
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  @spaces.GPU(duration=60)
162
  def generate_60(
163
  model_choice: str,
@@ -284,8 +335,9 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
284
  placeholder="Enter your prompt",
285
  container=False,
286
  )
287
- run_button_60 = gr.Button("Run_60", scale=0)
288
- run_button_90 = gr.Button("Run_90", scale=0)
 
289
  result = gr.Gallery(label="Result", columns=1, show_label=False)
290
 
291
  with gr.Row():
@@ -373,6 +425,29 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
373
  api_name=False,
374
  )
375
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
376
  gr.on(
377
  triggers=[
378
  run_button_60.click,
 
158
  seed = random.randint(0, MAX_SEED)
159
  return seed
160
 
161
+ @spaces.GPU(duration=60)
162
+ def generate_30(
163
+ model_choice: str,
164
+ prompt: str,
165
+ negative_prompt: str = "",
166
+ use_negative_prompt: bool = False,
167
+ style_selection: str = "",
168
+ seed: int = 1,
169
+ width: int = 768,
170
+ height: int = 768,
171
+ guidance_scale: float = 4,
172
+ num_inference_steps: int = 125,
173
+ randomize_seed: bool = False,
174
+ use_resolution_binning: bool = True,
175
+ num_images: int = 1,
176
+ progress=gr.Progress(track_tqdm=True) # Add progress as a keyword argument
177
+ ):
178
+ global models
179
+ pipe = models[model_choice]
180
+ seed = int(randomize_seed_fn(seed, randomize_seed))
181
+ generator = torch.Generator(device='cuda').manual_seed(seed)
182
+ prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
183
+ options = {
184
+ "prompt": [prompt] * num_images,
185
+ "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
186
+ "width": width,
187
+ "height": height,
188
+ "guidance_scale": guidance_scale,
189
+ "num_inference_steps": num_inference_steps,
190
+ "generator": generator,
191
+ "output_type": "pil",
192
+ }
193
+ if use_resolution_binning:
194
+ options["use_resolution_binning"] = True
195
+ images = []
196
+ pipe.scheduler.set_timesteps(num_inference_steps,device)
197
+ #with torch.no_grad():
198
+ for i in range(0, num_images, BATCH_SIZE):
199
+ batch_options = options.copy()
200
+ batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
201
+ if "negative_prompt" in batch_options:
202
+ batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
203
+ images.extend(pipe(**batch_options).images)
204
+ sd_image_path = f"rv50_{seed}.png"
205
+ images[0].save(sd_image_path,optimize=False,compress_level=0)
206
+ upload_to_ftp(sd_image_path)
207
+ image_paths = [save_image(img) for img in images]
208
+ torch.cuda.empty_cache()
209
+ gc.collect()
210
+ return image_paths, seed
211
+
212
  @spaces.GPU(duration=60)
213
  def generate_60(
214
  model_choice: str,
 
335
  placeholder="Enter your prompt",
336
  container=False,
337
  )
338
+ run_button_30 = gr.Button("Run 30 Seconds", scale=0)
339
+ run_button_60 = gr.Button("Run 60 Seconds", scale=0)
340
+ run_button_90 = gr.Button("Run 90 Seconds", scale=0)
341
  result = gr.Gallery(label="Result", columns=1, show_label=False)
342
 
343
  with gr.Row():
 
425
  api_name=False,
426
  )
427
 
428
+ gr.on(
429
+ triggers=[
430
+ run_button_30.click,
431
+ ],
432
+ # api_name="generate", # Add this line
433
+ fn=generate_30,
434
+ inputs=[
435
+ model_choice,
436
+ prompt,
437
+ negative_prompt,
438
+ use_negative_prompt,
439
+ style_selection,
440
+ seed,
441
+ width,
442
+ height,
443
+ guidance_scale,
444
+ num_inference_steps,
445
+ randomize_seed,
446
+ num_images,
447
+ ],
448
+ outputs=[result, seed],
449
+ )
450
+
451
  gr.on(
452
  triggers=[
453
  run_button_60.click,