src/aiunn/finetune.py

@@ -23,22 +23,25 @@ class ImageDataset(Dataset):
         row = self.dataframe.iloc[idx]
 
         try:
-            # Directly use bytes data for PNG images
-            low_res_bytes = row['image_512']
-            high_res_bytes = row['image_1024']
+            # Convert string data to bytes if needed
+            low_res_bytes = row['image_512'].encode() if isinstance(row['image_512'], str) else row['image_512']
+            high_res_bytes = row['image_1024'].encode() if isinstance(row['image_1024'], str) else row['image_1024']
             
-            # Create in-memory streams
+            # Create BytesIO objects
             low_res_stream = io.BytesIO(low_res_bytes)
             high_res_stream = io.BytesIO(high_res_bytes)
 
-            # Open images with explicit RGB conversion
+            # Open images
             low_res_image = Image.open(low_res_stream).convert('RGB')
             high_res_image = Image.open(high_res_stream).convert('RGB')
 
+            # Close the streams
+            low_res_stream.close()
+            high_res_stream.close()
+
         except Exception as e:
             raise ValueError(f"Image loading failed: {str(e)}")
 
-        # Apply transformations if specified
         if self.transform:
             low_res_image = self.transform(low_res_image)
             high_res_image = self.transform(high_res_image)