someone-in-the-world commited on
Commit
629c301
·
unverified ·
1 Parent(s): de8e11b

Stop Inductor compilation progress from leaking into the Gradio UI (#33)

Browse files

gr.Progress(track_tqdm=True) monkey-patches tqdm globally, so it was
picking up TorchInductor's own "Inductor Compilation" progress bars
(triggered by torch.compile on the upscale model, and by the quantized
transformer/text_encoder) regardless of their own disable flag. Drop
track_tqdm and report the generation stage's progress explicitly via
callback_on_step_end instead, matching the interpolation/upscale stages.

Files changed (1) hide show
  1. app.py +7 -2
app.py CHANGED
@@ -161,11 +161,15 @@ def run_inference(
161
  seed,
162
  frame_multiplier,
163
  upscale_output,
164
- progress=gr.Progress(track_tqdm=True),
165
  ):
166
  print_infer_start(prompt, negative_prompt, seed, steps, guidance_scale, frame_multiplier, upscale_output)
167
  t_start = _time.perf_counter()
168
 
 
 
 
 
169
  print_stage_start("generation")
170
  t0 = _time.perf_counter()
171
  try:
@@ -180,6 +184,7 @@ def run_inference(
180
  num_inference_steps=int(steps),
181
  generator=torch.Generator(device="cuda").manual_seed(seed),
182
  output_type="np",
 
183
  )
184
  except Exception as e:
185
  print_stage_error("generation", e)
@@ -246,7 +251,7 @@ def generate_video(
246
  randomize_seed,
247
  frame_multiplier,
248
  upscale_output,
249
- progress=gr.Progress(track_tqdm=True),
250
  ):
251
  if input_image is None:
252
  raise gr.Error("Please upload an input image.")
 
161
  seed,
162
  frame_multiplier,
163
  upscale_output,
164
+ progress=gr.Progress(),
165
  ):
166
  print_infer_start(prompt, negative_prompt, seed, steps, guidance_scale, frame_multiplier, upscale_output)
167
  t_start = _time.perf_counter()
168
 
169
+ def _report_generation_progress(pipe_, step_index, timestep, callback_kwargs):
170
+ progress(0.7 * (step_index + 1) / int(steps), desc=f"Generating ({step_index + 1}/{int(steps)} steps)...")
171
+ return callback_kwargs
172
+
173
  print_stage_start("generation")
174
  t0 = _time.perf_counter()
175
  try:
 
184
  num_inference_steps=int(steps),
185
  generator=torch.Generator(device="cuda").manual_seed(seed),
186
  output_type="np",
187
+ callback_on_step_end=_report_generation_progress,
188
  )
189
  except Exception as e:
190
  print_stage_error("generation", e)
 
251
  randomize_seed,
252
  frame_multiplier,
253
  upscale_output,
254
+ progress=gr.Progress(),
255
  ):
256
  if input_image is None:
257
  raise gr.Error("Please upload an input image.")