Upload inference.py

inference.py

@@ -64,14 +64,17 @@ def inference(pipe, img1, img2, num_steps):
     image = PIL.ImageOps.exif_transpose(image)
     
     all_images = []
+    
     def cb_fn(step, timestep, latents):
-        # latents(torch.Tensor) -> image via VAE decode
+        # 1) Décoder les latents -> DecoderOutput
         with torch.no_grad():
-            decoded = pipe.vae.decode(latents / pipe.vae.config.scaling_factor)
-        # post‐traitement en PIL
-        img = pipe.numpy_to_pil(decoded.cpu().clamp(0,1))[0]
+            decoded_output = pipe.vae.decode(latents / pipe.vae.config.scaling_factor)
+        # 2) Extraire le tenseur : .sample contient le batch de sorties
+        decoded_tensor = decoded_output.sample  # type: torch.Tensor
+        # 3) Passer sur CPU, clampler et convertir en PIL
+        img = pipe.numpy_to_pil(decoded_tensor.cpu().clamp(0, 1))[0]
         all_images.append(img)
-
+    
     num_inference_steps = num_steps
     image_guidance_scale = 1.9
     guidance_scale = 10