优化进度显示

app.py

@@ -93,6 +93,7 @@ def validate_dimensions(width: int, height: int) -> None:
 
 
 
+progress=gr.Progress()
 
 @spaces.GPU
 def generate(
@@ -106,17 +107,27 @@ def generate(
     seed: int,
     randomize_seed: bool,
     guidance_scale: float,
-    num_inference_steps: int,
-    progress:gr.Progress=gr.Progress(track_tqdm=True),
+    num_inference_steps: int
 ):
+    progress(0,desc="Starting")
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
     """Generate images based on the given parameters."""
-    start_time = time.time()
     upscaler_pipe = None
     backup_scheduler = None
 
+    def callback1(pipe, step, timestep, callback_kwargs):
+        progress_value = 0.1 + ((step+1.0)/num_inference_steps)*(0.5/1.0)
+        progress(progress_value, desc=f"Image generating, {step + 1}/{num_inference_steps} steps")
+        return callback_kwargs
+    
+    def callback2(pipe, step, timestep, callback_kwargs):
+        progress_value = 0.6 + ((step+1.0)/num_inference_steps)*(0.4/1.0)
+        progress(progress_value, desc=f"Image optimizing, {step + 1}/{num_inference_steps} steps")
+        return callback_kwargs
+
     try:
         # Memory management
         torch.cuda.empty_cache()
@@ -140,6 +151,7 @@ def generate(
 
         upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
 
+        progress(0.1,desc="Image generating")
         latents = pipe(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -149,8 +161,9 @@ def generate(
                 num_inference_steps=num_inference_steps,
                 generator=generator,
                 output_type="latent",
+                callback_on_step_end=callback1
             ).images
-
+        progress(0.6,desc="Image optimizing")
         upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
         images = upscaler_pipe(
                 prompt=prompt,
@@ -161,9 +174,11 @@ def generate(
                 strength=upscaler_strength,
                 generator=generator,
                 output_type="pil",
+                callback_on_step_end=callback2
             ).images
-        return images[0]
-    
+        out_img = images[0]
+        progress(1,desc="Complete")
+        return out_img
     except GenerationError as e:
         logger.warning(f"Generation validation error: {str(e)}")
         raise gr.Error(str(e))