Update app.py

app.py

@@ -13,13 +13,29 @@ from tqdm.auto import tqdm
 import random
 import gradio as gr
 
+
+import torchvision.transforms as transforms
+
+def tensor_to_pil(tensor):
+    # Assuming tensor is normalized
+    unnormalize = transforms.Normalize(mean=[-0.5 / 0.5, -0.5 / 0.5, -0.5 / 0.5], std=[1/0.5, 1/0.5, 1/0.5])
+    tensor = unnormalize(tensor)
+    tensor = tensor.clamp(0, 1)
+    tensor = tensor.permute(1, 2, 0)  # Convert from CxHxW to HxWxC
+    image = Image.fromarray((tensor.numpy() * 255).astype('uint8'))
+    return image
+
+
 def generate_images(prompt, guidance_scale, n_samples, num_inference_steps):
     seeds = [random.randint(1, 10000) for _ in range(n_samples)]
     images = []
     for seed in tqdm(seeds):
         torch.manual_seed(seed)
-        image = pipe(prompt, num_inference_steps=num_inference_steps,guidance_scale=guidance_scale).images[0]
-        images.append(image)
+        #tensor = pipe(prompt, num_inference_steps=num_inference_steps,guidance_scale=guidance_scale).images[0]
+        tensor_image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
+        # Convert tensor to PIL Image
+        pil_image = tensor_to_pil(tensor_image)
+        images.append(pil_image)        
     return images
 
 def gr_generate_images(prompt: str, num_images = 1, num_inference = 20, guidance_scale = 8 ):