Upload inference.py

inference.py

@@ -11,7 +11,7 @@ from diffusers import StableDiffusionInstructPix2PixPipeline, UNet2DModel, Autoe
 import numpy as np
 import torch.nn as nn
 from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
-
+from PIL import Image
 
 class UNetNoCondWrapper(nn.Module):
     def __init__(self, base_unet: UNet2DModel):
@@ -66,13 +66,17 @@ def inference(pipe, img1, img2, num_steps):
     all_images = []
     
     def cb_fn(step, timestep, latents):
-        # 1) Décoder les latents -> DecoderOutput
+        # 1) Décoder
         with torch.no_grad():
             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]
+        decoded_tensor = decoded_output.sample  # (B, C, H, W)
+        
+        # 2) Transformer en NumPy (channels last) et en uint8 [0–255]
+        t = decoded_tensor.cpu().clamp(0, 1)[0]           # (C, H, W)
+        arr = (t.permute(1, 2, 0).numpy() * 255).astype(np.uint8)  # (H, W, C)
+        
+        # 3) Créer la PIL.Image
+        img = Image.fromarray(arr)
         all_images.append(img)
     
     num_inference_steps = num_steps