Fix: full denoising teimsteps

app.py

@@ -104,7 +104,7 @@ def generate_images(selected_class, num_images, progress=gr.Progress()):
     x = torch.randn(num_images, 1, 28, 28).to(device)
     y = torch.tensor([class_idx] * num_images).to(device)
     
-    for t in progress.tqdm(range(2), desc="Generating image", total=noise_scheduler.config.num_train_timesteps): # 
+    for t in progress.tqdm(noise_scheduler.timesteps, desc="Generating image", total=noise_scheduler.config.num_train_timesteps): # 
         with torch.no_grad():
             residual = model_classcond_native(x, t, y)
         x = noise_scheduler.step(residual, t, x).prev_sample