Fix LLaVA fallback image loading in Stage 1

code/train_production.py

@@ -14,6 +14,7 @@ import json
 import math
 import os
 import time
+import zipfile
 from collections import defaultdict
 from dataclasses import dataclass
 from io import BytesIO
@@ -359,13 +360,20 @@ def tokenize_prompt_and_target(
 
 
 def preprocess_image_for_student(img: object, img_size: int) -> Tuple[torch.Tensor, Image.Image]:
-    if isinstance(img, str):
-        img = Image.open(img).convert("RGB")
-    elif isinstance(img, dict) and "bytes" in img:
-        img = Image.open(BytesIO(img["bytes"])).convert("RGB")
-    elif isinstance(img, Image.Image):
-        img = img.convert("RGB")
-    else:
+    try:
+        if isinstance(img, str):
+            img = Image.open(img).convert("RGB")
+        elif isinstance(img, dict) and "bytes" in img:
+            img = Image.open(BytesIO(img["bytes"])).convert("RGB")
+        elif isinstance(img, dict) and "zip_path" in img and "member" in img:
+            with zipfile.ZipFile(img["zip_path"], "r") as archive:
+                with archive.open(img["member"], "r") as member_file:
+                    img = Image.open(member_file).convert("RGB")
+        elif isinstance(img, Image.Image):
+            img = img.convert("RGB")
+        else:
+            img = Image.new("RGB", (img_size, img_size), (128, 128, 128))
+    except Exception:
         img = Image.new("RGB", (img_size, img_size), (128, 128, 128))
 
     pil_image = img
@@ -420,6 +428,7 @@ class NormalizedVisionLanguageDataset(Dataset):
 def build_llava_records(max_samples: Optional[int]) -> HFDataset:
     print("Loading LLaVA-Pretrain dataset...")
     dataset_root = None
+    images_zip_path = None
     try:
         data = load_dataset("liuhaotian/LLaVA-Pretrain", split="train")
     except Exception as exc:
@@ -430,6 +439,7 @@ def build_llava_records(max_samples: Optional[int]) -> HFDataset:
             allow_patterns=["blip_laion_cc_sbu_558k.json", "images.zip"],
         )
         json_path = os.path.join(dataset_root, "blip_laion_cc_sbu_558k.json")
+        images_zip_path = os.path.join(dataset_root, "images.zip")
         data = load_dataset("json", data_files={"train": json_path}, split="train")
     if max_samples:
         data = data.select(range(min(max_samples, len(data))))
@@ -455,7 +465,14 @@ def build_llava_records(max_samples: Optional[int]) -> HFDataset:
                 os.path.join(dataset_root, image_obj),
                 os.path.join(dataset_root, "images", image_obj),
             ]
-            image_obj = next((path for path in candidate_paths if os.path.exists(path)), image_obj)
+            resolved_path = next((path for path in candidate_paths if os.path.exists(path)), None)
+            if resolved_path:
+                image_obj = resolved_path
+            elif images_zip_path and os.path.exists(images_zip_path):
+                image_obj = {
+                    "zip_path": images_zip_path,
+                    "member": image_obj,
+                }
 
         return {
             "image": image_obj,