Fix image preview: build Parquet directly with PyArrow + HF schema metadata

datasets.Dataset was not writing the Arrow Image extension type metadata
correctly. Build the table directly with PyArrow using the exact struct
schema {bytes: binary, path: utf8} and a 'huggingface' schema metadata
key — the format the HF dataset viewer requires to render image columns.
Use replace_schema_metadata after concat to preserve the metadata.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -192,24 +192,43 @@ def log_inference(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
     if not HF_TOKEN or not DATASET_REPO:
         return
     try:
-        import tempfile
-        import datasets as ds_lib
+        import tempfile, json as _json
+        import pyarrow as pa
+        import pyarrow.parquet as pq
         from huggingface_hub import HfApi, hf_hub_download
 
-        features = ds_lib.Features({
-            "timestamp":        ds_lib.Value("string"),
-            "prompt":           ds_lib.Value("string"),
-            "seed":             ds_lib.Value("int32"),
-            "steps":            ds_lib.Value("int32"),
-            "guidance_scale":   ds_lib.Value("float32"),
-            "input_images":     ds_lib.Sequence(ds_lib.Image()),
-            "output_image":     ds_lib.Image(),
-            "duration_seconds": ds_lib.Value("float32"),
-            "input_width":      ds_lib.Value("int32"),
-            "input_height":     ds_lib.Value("int32"),
-            "success":          ds_lib.Value("bool"),
-            "error_message":    ds_lib.Value("string"),
-        })
+        # Image columns need Arrow struct {bytes: binary, path: utf8} plus
+        # a 'huggingface' schema metadata key for the HF viewer to render them.
+        img_struct = pa.struct([("bytes", pa.binary()), ("path", pa.string())])
+        hf_meta = _json.dumps({"info": {"features": {
+            "timestamp":        {"dtype": "string",  "_type": "Value"},
+            "prompt":           {"dtype": "string",  "_type": "Value"},
+            "seed":             {"dtype": "int32",   "_type": "Value"},
+            "steps":            {"dtype": "int32",   "_type": "Value"},
+            "guidance_scale":   {"dtype": "float32", "_type": "Value"},
+            "input_images":     {"feature": {"_type": "Image"}, "_type": "Sequence"},
+            "output_image":     {"_type": "Image"},
+            "duration_seconds": {"dtype": "float32", "_type": "Value"},
+            "input_width":      {"dtype": "int32",   "_type": "Value"},
+            "input_height":     {"dtype": "int32",   "_type": "Value"},
+            "success":          {"dtype": "bool",    "_type": "Value"},
+            "error_message":    {"dtype": "string",  "_type": "Value"},
+        }}}).encode()
+        schema = pa.schema([
+            ("timestamp",        pa.string()),
+            ("prompt",           pa.string()),
+            ("seed",             pa.int32()),
+            ("steps",            pa.int32()),
+            ("guidance_scale",   pa.float32()),
+            ("input_images",     pa.list_(img_struct)),
+            ("output_image",     img_struct),
+            ("duration_seconds", pa.float32()),
+            ("input_width",      pa.int32()),
+            ("input_height",     pa.int32()),
+            ("success",          pa.bool_()),
+            ("error_message",    pa.string()),
+        ], metadata={b"huggingface": hf_meta})
+
         def _to_jpeg(img, max_dim=1536, quality=85):
             if img is None:
                 return None
@@ -221,58 +240,58 @@ def log_inference(pil_inputs, output_pil, prompt, seed, steps, guidance_scale,
             img.convert("RGB").save(buf, format="JPEG", quality=quality)
             return buf.getvalue()
 
-        row = {
-            "timestamp":        datetime.now(timezone.utc).isoformat(),
-            "prompt":           prompt,
-            "seed":             int(seed),
-            "steps":            int(steps),
-            "guidance_scale":   float(guidance_scale),
-            "input_images":     [_to_jpeg(img) for img in pil_inputs],
-            "output_image":     _to_jpeg(output_pil),
-            "duration_seconds": float(duration_seconds),
-            "input_width":      int(input_width),
-            "input_height":     int(input_height),
-            "success":          bool(success),
-            "error_message":    str(error_message),
-        }
-        print(f"[log] building dataset row (success={success})")
-        new_ds = ds_lib.Dataset.from_dict(
-            {k: [v] for k, v in row.items()}, features=features
-        )
-        print(f"[log] row built — features: { {k: str(v) for k, v in new_ds.features.items()} }")
+        def _img(b):
+            return {"bytes": b, "path": None}
+
+        input_jpegs = [_to_jpeg(img) for img in pil_inputs]
+        output_jpeg = _to_jpeg(output_pil)
+
+        new_table = pa.table({
+            "timestamp":        pa.array([datetime.now(timezone.utc).isoformat()], type=pa.string()),
+            "prompt":           pa.array([prompt],              type=pa.string()),
+            "seed":             pa.array([int(seed)],           type=pa.int32()),
+            "steps":            pa.array([int(steps)],          type=pa.int32()),
+            "guidance_scale":   pa.array([float(guidance_scale)], type=pa.float32()),
+            "input_images":     pa.array([[_img(b) for b in input_jpegs]], type=pa.list_(img_struct)),
+            "output_image":     pa.array([_img(output_jpeg) if output_jpeg else None], type=img_struct),
+            "duration_seconds": pa.array([float(duration_seconds)], type=pa.float32()),
+            "input_width":      pa.array([int(input_width)],   type=pa.int32()),
+            "input_height":     pa.array([int(input_height)],  type=pa.int32()),
+            "success":          pa.array([bool(success)],      type=pa.bool_()),
+            "error_message":    pa.array([str(error_message)], type=pa.string()),
+        }, schema=schema)
+        print(f"[log] built row — success={success}, inputs={len(input_jpegs)}")
 
         today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
         path_in_repo = f"data/{today}.parquet"
         api = HfApi(token=HF_TOKEN)
-
-        print(f"[log] ensuring repo {DATASET_REPO} exists")
         api.create_repo(repo_id=DATASET_REPO, repo_type="dataset", private=True, exist_ok=True)
-        print(f"[log] repo ready")
 
         try:
-            print(f"[log] downloading existing {path_in_repo}")
             local_path = hf_hub_download(
                 repo_id=DATASET_REPO, filename=path_in_repo,
                 repo_type="dataset", token=HF_TOKEN,
             )
-            existing_ds = ds_lib.Dataset.from_parquet(local_path, features=features)
-            print(f"[log] existing file has {len(existing_ds)} row(s), appending")
-            combined_ds = ds_lib.concatenate_datasets([existing_ds, new_ds])
+            existing = pq.read_table(local_path)
+            combined = pa.concat_tables([existing, new_table])
+            combined = combined.replace_schema_metadata(schema.metadata)
+            print(f"[log] appending to existing {existing.num_rows} row(s)")
         except Exception as dl_err:
             print(f"[log] no existing file ({dl_err}), starting fresh")
-            combined_ds = new_ds
+            combined = new_table
 
-        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp_out:
-            tmp_out_path = tmp_out.name
-        combined_ds.to_parquet(tmp_out_path)
-        print(f"[log] uploading {path_in_repo} ({len(combined_ds)} row(s) total, {os.path.getsize(tmp_out_path)//1024}KB)")
+        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
+            tmp_path = tmp.name
+        pq.write_table(combined, tmp_path)
+        print(f"[log] uploading {path_in_repo} ({combined.num_rows} row(s), {os.path.getsize(tmp_path)//1024}KB)")
         api.upload_file(
-            path_or_fileobj=tmp_out_path, path_in_repo=path_in_repo,
+            path_or_fileobj=tmp_path, path_in_repo=path_in_repo,
             repo_id=DATASET_REPO, repo_type="dataset",
         )
         print(f"[log] upload done — {DATASET_REPO}/{path_in_repo}")
     except Exception as log_err:
-        print(f"[log] WARNING: failed to push log row: {log_err}")
+        import traceback as _tb
+        print(f"[log] WARNING: {log_err}\n{_tb.format_exc()}")
 
 
 @spaces.GPU

requirements.txt

@@ -3,7 +3,7 @@ git+https://github.com/huggingface/diffusers.git
 git+https://github.com/huggingface/peft.git
 transformers==4.57.1
 huggingface_hub
-datasets
+pyarrow
 sentencepiece
 torchvision
 kernels