revert

app.py

@@ -100,21 +100,20 @@ def validate_dimensions(width, height):
     return True, None
 
 @spaces.GPU()
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress()):
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     try:
-        # Use progress.update() instead of directly calling progress()
-        progress.update(0)
+        progress(0, desc="Starting generation...")
         
         # Validate that prompt is not empty
         if not prompt or prompt.strip() == "":
             return None, "Please provide a valid prompt."
 
-        progress.update(0.2, desc="Loading model...")
+        progress(0.1, desc="Loading image generation model...")
         pipe = get_image_gen_pipeline()
         if pipe is None:
             return None, "Image generation model is unavailable."
         
-        progress.update(0.4, desc="Validating dimensions...")
+        progress(0.2, desc="Validating dimensions...")
         is_valid, error_msg = validate_dimensions(width, height)
         if not is_valid:
             return None, error_msg
@@ -122,9 +121,10 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         
-        progress.update(0.6, desc="Generating image...")
+        progress(0.3, desc="Setting up generator...")
         generator = torch.Generator().manual_seed(seed)
         
+        progress(0.4, desc="Generating image...")
         with torch.autocast('cuda'):
             image = pipe(
                 prompt=prompt, 
@@ -137,7 +137,7 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
             ).images[0]
 
         torch.cuda.empty_cache()  # Clean up GPU memory after generation
-        progress.update(1.0, desc="Done!")
+        progress(1.0, desc="Done!")
         return image, seed
     except Exception as e:
         return None, f"Error generating image: {str(e)}"