Fix download bug

Author: coyotte508

packages/hub/src/utils/XetBlob.spec.ts

@@ -93,13 +93,13 @@ describe("XetBlob", () => {
 			"https://huggingface.co/celinah/xet-experiments/resolve/main/model5GB.safetensors",
 			{
 				headers: {
-					Range: "bytes=10000000-10100000",
+					Range: "bytes=10000000-10099999",
 				},
 			}
 		).then((res) => res.arrayBuffer());
 
 		expect(new Uint8Array(xetDownload)).toEqual(new Uint8Array(bridgeDownload));
-	});
+	}, 30_000);
 
 	it("should load text correctly when offset_into_range starts in a chunk further than the first", async () => {
 		const blob = new XetBlob({
@@ -147,8 +147,9 @@ describe("XetBlob", () => {
 			},
 		}).then((res) => res.text());
 
+		console.log("xet", text.length, "bridge", bridgeDownload.length);
 		expect(text.length).toBe(bridgeDownload.length);
-	});
+	}, 30_000);
 
 	describe("bg4_regoup_bytes", () => {
 		it("should regroup bytes when the array is %4 length", () => {

packages/hub/src/utils/XetBlob.ts

@@ -183,10 +183,10 @@ export class XetBlob extends Blob {
 			let readBytesToSkip = reconstructionInfo.offset_into_first_range;
 
 			for (const term of reconstructionInfo.terms) {
-				// console.log("term", term, readBytesToSkip, totalBytesRead);
 				if (totalBytesRead >= maxBytes) {
 					break;
 				}
+
 				const fetchInfo = reconstructionInfo.fetch_info[term.hash].find(
 					(info) => info.range.start <= term.range.start && info.range.end >= term.range.end
 				);
@@ -240,6 +240,7 @@ export class XetBlob extends Blob {
 									continue readChunks;
 								}
 								result.value = result.value.slice(bytesToSkip);
+								bytesToSkip = 0;
 							}
 							if (leftoverBytes) {
 								result.value = new Uint8Array([...leftoverBytes, ...result.value]);
@@ -314,21 +315,26 @@ export class XetBlob extends Blob {
 									    )
 									  : result.value.slice(0, chunkHeader.compressed_length);
 
+							let bytesToYield: Uint8Array;
 							if (readBytesToSkip) {
-								yield uncompressed.slice(
-									readBytesToSkip,
-									Math.min(uncompressed.length, readBytesToSkip + maxBytes - totalBytesRead)
-								);
-								totalBytesRead += Math.min(uncompressed.length, readBytesToSkip + maxBytes - totalBytesRead);
+								const remainingBytes = Math.min(uncompressed.length - readBytesToSkip, maxBytes - totalBytesRead);
+								bytesToYield = uncompressed.slice(readBytesToSkip, readBytesToSkip + remainingBytes);
 								readBytesToSkip = 0;
 							} else {
-								yield uncompressed.slice(0, Math.min(uncompressed.length, maxBytes - totalBytesRead));
-								totalBytesRead += Math.min(uncompressed.length, maxBytes - totalBytesRead);
+								bytesToYield = uncompressed.slice(0, Math.min(uncompressed.length, maxBytes - totalBytesRead));
 							}
+
+							totalBytesRead += bytesToYield.length;
+							yield bytesToYield;
 							chunksToRead--;
+
+							result.value = result.value.slice(chunkHeader.compressed_length);
 						}
 					}
 				}
+
+				// Release the reader
+				await reader.cancel();
 			}
 		}