Update app.py

app.py

@@ -88,7 +88,7 @@ def load_checkpoint_for_inference(filepath, model_class):
 
 
 @spaces.GPU # Make this function run on GPU
-def generate_images(selected_class_name, num_samples):
+def generate_images(selected_class_name, num_samples, progress=gr.Progress()):
     print(f"Generating {num_samples} samples for class: {selected_class_name}")
 
     # Map class name to class ID
@@ -105,14 +105,14 @@ def generate_images(selected_class_name, num_samples):
     y = torch.full((num_samples,), label, dtype=torch.long).to(device)
 
     # Sampling loop
-    for i, t in tqdm(enumerate(noise_scheduler.timesteps)):
+    for i, t in progress.tqdm(enumerate(noise_scheduler.timesteps), total=len(noise_scheduler.timesteps), desc=f"Generating {selected_class_name} images"):
 
         # Get model pred
         with torch.no_grad():
             residual = model(x, t, y)  # Note that we pass in our label
 
         # Update sample with step
-        x = noise_scheduler.step(residual, t, x).prev_sample # Correctly update x
+        x = noise_scheduler.step(residual, t, x).prev_sample
 
     generated_pil_images = []
     for j in range(num_samples):