VinayHajare commited on
Commit
fa96828
1 Parent(s): 28b37f8

Update inference.py

Browse files
Files changed (1) hide show
  1. inference.py +13 -13
inference.py CHANGED
@@ -82,16 +82,16 @@ def generate_image_from_text_with_persistent_storage(caption, batch_size=4):
82
  # generate the images
83
  generated_images = generator(noise, sent_emb, eval=True).float()
84
 
85
- # Create a temporary directory
86
- with tempfile.TemporaryDirectory() as temp_dir:
87
- pil_images = []
88
- for idx, image_tensor in enumerate(generated_images.unbind(0)):
89
- # Save the image tensor to a temporary file
90
- image_path = os.path.join(temp_dir, f"image_{idx}.png")
91
- torchvision.utils.save_image(image_tensor.data, image_path, value_range=(-1, 1), normalize=True)
92
-
93
- # Load the saved image using Pillow
94
- pil_image = Image.open(image_path)
95
- pil_images.append(pil_image)
96
-
97
- return pil_images
 
82
  # generate the images
83
  generated_images = generator(noise, sent_emb, eval=True).float()
84
 
85
+ # Create a permanent directory if it doesn't exist
86
+ permanent_dir = "generated_images"
87
+ if not os.path.exists(permanent_dir):
88
+ os.makedirs(permanent_dir)
89
+
90
+ image_paths = []
91
+ for idx, image_tensor in enumerate(generated_images.unbind(0)):
92
+ # Save the image tensor to a permanent file
93
+ image_path = os.path.join(permanent_dir, f"image_{idx}.png")
94
+ torchvision.utils.save_image(image_tensor.data, image_path, value_range=(-1, 1), normalize=True)
95
+ image_paths.append(image_path)
96
+
97
+ return image_paths