Fix the diffussion model code

app.py

@@ -90,7 +90,7 @@ def generate_images(prompt, concept):
     
     for idx, loss_type in enumerate(loss_functions):
         try:
-            # Add detailed progress reporting
+            print(f"\n[{loss_type.upper()}] Starting image generation...")
             progress(idx/len(loss_functions), f"Starting {loss_type} image generation...")
             
             # Better memory management
@@ -149,7 +149,13 @@ def generate_images(prompt, concept):
             latents = latents * scheduler.init_noise_sigma
             
             # Diffusion process
+            total_steps = len(scheduler.timesteps)
             for i, t in enumerate(scheduler.timesteps):
+                current_progress = (idx + (i / total_steps)) / len(loss_functions)
+                progress_msg = f"[{loss_type.upper()}] Step {i+1}/{total_steps} ({(i+1)/total_steps*100:.1f}%)"
+                print(progress_msg)
+                progress(current_progress, progress_msg)
+                
                 latent_model_input = torch.cat([latents] * 2)
                 sigma = scheduler.sigmas[i]
                 latent_model_input = scheduler.scale_model_input(latent_model_input, t)
@@ -177,33 +183,21 @@ def generate_images(prompt, concept):
                         denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
                         denoised_images = denoised_images.requires_grad_()  # Enable gradients for images
                         loss = image_loss(denoised_images, loss_type, device, elastic_transformer)
-                        cond_grad = torch.autograd.grad(loss * loss_scale, latents)[0]
+                        # Ensure latents_x0 requires grad
+                        if not latents_x0.requires_grad:
+                            latents_x0 = latents_x0.requires_grad_()
+                        cond_grad = torch.autograd.grad(loss * loss_scale, latents_x0)[0]
                     
                     latents = latents.detach() - cond_grad * sigma**2
-                
-                # Diffusion process with progress updates
-                for i, t in enumerate(scheduler.timesteps):
-                    current_progress = (idx + (i / len(scheduler.timesteps))) / len(loss_functions)
-                    progress(current_progress, f"Generating {loss_type} image: Step {i+1}/{len(scheduler.timesteps)}")
-                    
-                    # Apply loss less frequently for speed
-                    if loss_type != 'none' and i % 8 == 0:  # Changed from 5 to 8
-                        with torch.set_grad_enabled(True):
-                            # Enable gradients for images
-                            denoised_images = pipe.vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5
-                            denoised_images = denoised_images.requires_grad_()  # Enable gradients for images
-                            loss = image_loss(denoised_images, loss_type, device, elastic_transformer)
-                            cond_grad = torch.autograd.grad(loss * loss_scale, latents)[0]
-                        
-                        latents = latents.detach() - cond_grad * sigma**2
-                
+
                 latents = scheduler.step(noise_pred, t, latents).prev_sample
                 
                 # Clear CUDA cache more efficiently
                 if torch.cuda.is_available() and i % 10 == 0:
                     torch.cuda.empty_cache()
             
-                progress(idx/len(loss_functions), f"Finalizing {loss_type} image...")
+            # Remove the nested diffusion loop and move finalization outside
+            progress(idx/len(loss_functions), f"Finalizing {loss_type} image...")
             
             # Proper latent to image conversion
             latents = (1 / 0.18215) * latents
@@ -220,12 +214,13 @@ def generate_images(prompt, concept):
             
         except Exception as e:
             print(f"Error generating {loss_type} image: {e}")
-            continue  # Continue to next loss type instead of returning None
+            continue
     
-    # At the end of the function
+    # At the end of the function, outside the loop
     try:
         if len(all_images) == 0:
             raise Exception("No images were generated successfully")
+        print("\nAll images generated successfully!")
         return [img for img, _ in all_images]
     except Exception as e:
         print(f"Error in generate_images: {e}")