File size: 39,216 Bytes
16d51ab
 
 
 
833555d
16d51ab
 
833555d
 
 
 
 
 
 
 
 
16d51ab
 
 
 
 
 
833555d
16d51ab
 
833555d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16d51ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
833555d
16d51ab
 
 
 
833555d
16d51ab
 
 
 
 
 
 
 
 
 
 
 
 
 
833555d
16d51ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49cbc74
16d51ab
 
 
 
 
833555d
16d51ab
 
833555d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
16d51ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c9cc09d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/Qwen-Image-Edit-Angles\n"
     ]
    }
   ],
   "source": [
    "%cd /home/ubuntu/Qwen-Image-Edit-Angles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b65a5e8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3/dist-packages/sklearn/utils/fixes.py:25: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  from pkg_resources import parse_version  # type: ignore\n",
      "2025-11-24 16:55:55.657889: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-11-24 16:55:55.671869: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "E0000 00:00:1764003355.688913 3244532 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "E0000 00:00:1764003355.694358 3244532 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "W0000 00:00:1764003355.707749 3244532 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1764003355.707764 3244532 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1764003355.707767 3244532 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "W0000 00:00:1764003355.707768 3244532 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
      "2025-11-24 16:55:55.712504: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.local/lib/python3.10/site-packages/google/api_core/_python_version_support.py:266: FutureWarning: You are using a Python version (3.10.12) which Google will stop supporting in new releases of google.api_core once it reaches its end of life (2026-10-04). Please upgrade to the latest Python version, or at least Python 3.11, to continue receiving updates for google.api_core past that date.\n",
      "  warnings.warn(message, FutureWarning)\n"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
     ]
    },
    {
     "ename": "AttributeError",
     "evalue": "'MessageFactory' object has no attribute 'GetPrototype'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;31mAttributeError\u001b[0m: 'MessageFactory' object has no attribute 'GetPrototype'"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Skipping import of cpp extensions due to incompatible torch version 2.9.1+cu128 for torchao version 0.14.1             Please see https://github.com/pytorch/ao/issues/2919 for more info\n",
      "TMA benchmarks will be running without grid constant TMA descriptor.\n",
      "WARNING:bitsandbytes.cextension:Could not find the bitsandbytes CUDA binary at PosixPath('/usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cuda128.so')\n",
      "ERROR:bitsandbytes.cextension:Could not load bitsandbytes native library: /lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.32' not found (required by /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cpu.so)\n",
      "Traceback (most recent call last):\n",
      "  File \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py\", line 85, in <module>\n",
      "    lib = get_native_library()\n",
      "  File \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cextension.py\", line 72, in get_native_library\n",
      "    dll = ct.cdll.LoadLibrary(str(binary_path))\n",
      "  File \"/usr/lib/python3.10/ctypes/__init__.py\", line 452, in LoadLibrary\n",
      "    return self._dlltype(name)\n",
      "  File \"/usr/lib/python3.10/ctypes/__init__.py\", line 374, in __init__\n",
      "    self._handle = _dlopen(self._name, mode)\n",
      "OSError: /lib/x86_64-linux-gnu/libstdc++.so.6: version `GLIBCXX_3.4.32' not found (required by /usr/local/lib/python3.10/dist-packages/bitsandbytes/libbitsandbytes_cpu.so)\n",
      "WARNING:bitsandbytes.cextension:\n",
      "CUDA Setup failed despite CUDA being available. Please run the following command to get more information:\n",
      "\n",
      "python -m bitsandbytes\n",
      "\n",
      "Inspect the output of the command and see if you can locate CUDA libraries. You might need to add them\n",
      "to your LD_LIBRARY_PATH. If you suspect a bug, please take the information from python -m bitsandbytes\n",
      "and open an issue at: https://github.com/bitsandbytes-foundation/bitsandbytes/issues\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "599e401d77bc49edaaacbfe6f55032cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 7 files:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import huggingface_hub \n",
    "from qwenimage.datamodels import QwenConfig\n",
    "from qwenimage.foundation import QwenImageFoundationSaveInterm\n",
    "from datasets import concatenate_datasets, load_dataset, interleave_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab2ef566",
   "metadata": {},
   "outputs": [],
   "source": [
    "repo_tree = huggingface_hub.list_repo_tree(\n",
    "    \"WeiChow/CrispEdit-2M\",\n",
    "    \"data\",\n",
    "    repo_type=\"dataset\",\n",
    ")\n",
    "\n",
    "all_paths = []\n",
    "for i in repo_tree:\n",
    "    all_paths.append(i.path)\n",
    "\n",
    "parquet_prefixes = set()\n",
    "for path in all_paths:\n",
    "    if path.endswith('.parquet'):\n",
    "        filename = path.split('/')[-1]\n",
    "        if '_' in filename:\n",
    "            prefix = filename.split('_')[0]\n",
    "            parquet_prefixes.add(prefix)\n",
    "\n",
    "print(\"Found parquet prefixes:\", sorted(parquet_prefixes))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5505f47d",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_per = 1\n",
    "\n",
    "EDIT_TYPES = [\n",
    "    \"color\",\n",
    "    \"style\",\n",
    "    \"replace\",\n",
    "    \"remove\",\n",
    "    \"add\",\n",
    "    \"motion change\",\n",
    "    \"background change\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cdf51dbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "all_edit_datasets = []\n",
    "for edit_type in EDIT_TYPES:\n",
    "    to_concat = []\n",
    "    for ds_n in range(total_per):\n",
    "        ds = load_dataset(\"parquet\", data_files=f\"/data/CrispEdit/{edit_type}_{ds_n:05d}.parquet\", split=\"train\")\n",
    "        to_concat.append(ds)\n",
    "    edit_type_concat = concatenate_datasets(to_concat)\n",
    "    all_edit_datasets.append(edit_type_concat)\n",
    "\n",
    "# consistent ordering for indexing, also allow extension by increasing total_per\n",
    "join_ds = interleave_datasets(all_edit_datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0c59b9d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Turn dress positioned in the lower central area into red with floral patterns\n",
      "Please recreate this in digital painting artistic style\n",
      "replace the largemouth bass with a trout\n",
      "remove the four grey extraterrestrials observing the scene\n",
      "Add the lid back onto the jar of white cream\n",
      "The character stops exhaling smoke and closes their mouth.\n",
      "change the background to a mystical forest\n",
      "Turn cap positioned in the upper-central area into red\n",
      "Draw this image using fantasy art technique\n",
      "replace the alien spaceship with a futuristic submarine\n",
      "remove the morbidly obese Frankenstein monster from the school bus\n",
      "add a halved lemon\n",
      "No change in the subject's action, pose, or facial expression.\n",
      "change the background to a serene starry night sky\n",
      "Turn baby positioned in the central area into darker skin tone\n",
      "Convert this image to graffiti artwork\n",
      "replace the pregnant woman with an elderly woman holding a walking stick\n",
      "remove the Wakandan technology elements from the townscape\n",
      "add a large, serene pool area\n",
      "The woman lowers her hand and straightens her posture.\n",
      "change the background to a serene beach at sunset\n",
      "Turn black tires positioned in the lower left area into white\n",
      "Reimagine this image in low poly artistic style\n",
      "replace the human skull with a flower crown\n",
      "remove the pile of cheap broken cars from the white background\n",
      "Add three plastic eggs in blue, green, and yellow.\n",
      "The person holds a small piece of food in one hand while still holding the phone in the other.\n",
      "change the background to a futuristic cityscape\n",
      "Turn vehicles positioned in the lower-left area into polished\n",
      "Turn this photo into oil painting artwork\n",
      "replace the spike two-handed mace with a sword\n",
      "remove the baby version of Batman from the image\n",
      "add sparkling rhinestones to the bouquet of white roses\n",
      "A person performs a pull-up with knees raised and holding a medicine ball.\n",
      "change the background to a bustling city street at night\n",
      "Turn grass positioned in the bottom area into dry\n",
      "Transform this image using polaroid artistic approach\n",
      "replace the plant factory with a castle\n",
      "remove the palace on the moon\n",
      "Add the classic blue muscle car back onto the paved road\n",
      "The subject changes from having hands clasped behind their back to holding their arms out to the sides.\n",
      "change the background to a serene meadow under a clear blue sky\n",
      "Turn wheel positioned in the central area into black with red brake caliper\n",
      "Render this image as vintage photography art\n",
      "replace the horned frog with a deer\n",
      "erase the abstract body featuring many limbs\n",
      "add the golden fork to the white lace plate\n",
      "The person turns their head to the side.\n",
      "change the background to a starry night sky with glowing constellations\n",
      "Turn flowers positioned in the lower-central area into pink\n",
      "Can you render this image as art nouveau art?\n",
      "replace the monoliths with ancient stone statues\n",
      "remove the colorful fireworks in the sky\n",
      "Add a green bottle near the cupcakes\n",
      "The man's hand gesture changes from an open motion to a partially closed motion.\n",
      "change the background to a futuristic cityscape\n",
      "Turn chest pack positioned in the upper-central area into metallic silver\n",
      "Transform this image using glitch art artistic approach\n",
      "replace the fiery heart with a glowing moon\n",
      "remove the purple rally stripes from the military aircraft\n",
      "Add a hat to the man\n",
      "The person lowers their right hand from a raised position to resting on their lap while keeping their left hand holding a phone.\n",
      "change the background to a futuristic cityscape\n",
      "Turn tables positioned in the lower central area into wooden\n",
      "Illustrate this in woodcut format\n",
      "replace the knight with a wizard\n",
      "remove the meticulously designed object in the center\n",
      "add a golden retriever lying beside the woman\n",
      "A woman raises her left hand slightly higher and moves her right hand outward.\n",
      "change the background to a city skyline at night\n",
      "Turn frog positioned in the upper-central area into dark green\n",
      "Reimagine this image in graffiti artistic style\n",
      "replace the snowman with a sandcastle\n",
      "remove the FABLER logo from the illustration\n",
      "add the colorful flower mural behind the four friends\n",
      "The man turns his head slightly to his left.\n",
      "change the background to a starry night sky\n",
      "Turn fence positioned in the lower central area into wooden\n",
      "Could you convert this to pop art artwork?\n",
      "replace the campfire with a small table\n",
      "remove the fairies from the forest\n",
      "add a hand stirring the glass pitcher\n",
      "The person adjusts the object in their hands, moving from holding it with both hands to manipulating it with one hand.\n",
      "change the background to a calm, serene ocean at sunset\n",
      "Turn house positioned in the right-central area into wooden\n",
      "Convert this image to low poly artwork\n",
      "replace the bird with a butterfly\n",
      "remove the engraved gem from the image\n",
      "Add lounge chairs and umbrellas by the pool.\n",
      "A woman reaches to remove the lid from a kitchen appliance.\n",
      "change the background to a serene blue sky with fluffy white clouds\n",
      "Turn piano positioned in the lower-left area into white\n",
      "Please transform this image into isometric style\n",
      "replace the glasses with a monocle\n",
      "remove the circus carnival background\n",
      "Add people relaxing on the grass\n",
      "A man adjusts the virtual reality headset on his head.\n",
      "change the background to a stormy sea with a pirate ship in the distance\n",
      "Turn buildings positioned in the upper-right area into brightly colored structures\n",
      "Please recreate this in street art artistic style\n",
      "replace the father with a teacher\n",
      "remove the black orb hovering in the galaxy scene\n",
      "add back the figure who appears to be asleep\n",
      "The person moves the camera from one hand to both hands, holding it more securely.\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_3244532/4022048012.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0md\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mjoin_ds\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"instruction\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/arrow_dataset.py\u001b[0m in \u001b[0;36m__iter__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   2491\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2492\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_rows\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2493\u001b[0;31m                 yield self._getitem(\n\u001b[0m\u001b[1;32m   2494\u001b[0m                     \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2495\u001b[0m                 )\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/arrow_dataset.py\u001b[0m in \u001b[0;36m_getitem\u001b[0;34m(self, key, **kwargs)\u001b[0m\n\u001b[1;32m   2856\u001b[0m         \u001b[0mformatter\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_formatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mformat_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfeatures\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_info\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mformat_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2857\u001b[0m         \u001b[0mpa_subtable\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mquery_table\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindices\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_indices\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2858\u001b[0;31m         formatted_output = format_table(\n\u001b[0m\u001b[1;32m   2859\u001b[0m             \u001b[0mpa_subtable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mformatter\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mformat_columns\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mformat_columns\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput_all_columns\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moutput_all_columns\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2860\u001b[0m         )\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/formatting/formatting.py\u001b[0m in \u001b[0;36mformat_table\u001b[0;34m(table, key, formatter, format_columns, output_all_columns)\u001b[0m\n\u001b[1;32m    664\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    665\u001b[0m         \u001b[0mpa_table_to_format\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpa_table\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdrop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcol\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mcol\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mpa_table\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolumn_names\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mcol\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mformat_columns\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 666\u001b[0;31m         \u001b[0mformatted_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mformatter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpa_table_to_format\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquery_type\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mquery_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    667\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0moutput_all_columns\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    668\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mformatted_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mMutableMapping\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/formatting/formatting.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, pa_table, query_type)\u001b[0m\n\u001b[1;32m    409\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpa_table\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mpa\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mquery_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mRowFormat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mColumnFormat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mBatchFormat\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    410\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mquery_type\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"row\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 411\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat_row\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpa_table\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    412\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mquery_type\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"column\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    413\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat_column\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpa_table\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/formatting/formatting.py\u001b[0m in \u001b[0;36mformat_row\u001b[0;34m(self, pa_table)\u001b[0m\n\u001b[1;32m    458\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mLazyRow\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpa_table\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m         \u001b[0mrow\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpython_arrow_extractor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextract_row\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpa_table\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 460\u001b[0;31m         \u001b[0mrow\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpython_features_decoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode_row\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrow\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    461\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/formatting/formatting.py\u001b[0m in \u001b[0;36mdecode_row\u001b[0;34m(self, row)\u001b[0m\n\u001b[1;32m    222\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    223\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mdecode_row\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mdict\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 224\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode_example\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrow\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfeatures\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mrow\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    225\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    226\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mdecode_column\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcolumn\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcolumn_name\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/features/features.py\u001b[0m in \u001b[0;36mdecode_example\u001b[0;34m(self, example, token_per_repo_id)\u001b[0m\n\u001b[1;32m   2103\u001b[0m         \"\"\"\n\u001b[1;32m   2104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2105\u001b[0;31m         return {\n\u001b[0m\u001b[1;32m   2106\u001b[0m             \u001b[0mcolumn_name\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mdecode_nested_example\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeature\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2107\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_column_requires_decoding\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcolumn_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/features/features.py\u001b[0m in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m   2104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2105\u001b[0m         return {\n\u001b[0;32m-> 2106\u001b[0;31m             \u001b[0mcolumn_name\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mdecode_nested_example\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeature\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2107\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_column_requires_decoding\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcolumn_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2108\u001b[0m             \u001b[0;32melse\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/features/features.py\u001b[0m in \u001b[0;36mdecode_nested_example\u001b[0;34m(schema, obj, token_per_repo_id)\u001b[0m\n\u001b[1;32m   1412\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mschema\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"decode_example\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mschema\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"decode\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1413\u001b[0m         \u001b[0;31m# we pass the token to read and decode files from private repositories in streaming mode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1414\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mschema\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode_example\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtoken_per_repo_id\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mobj\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1415\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1416\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/datasets/features/image.py\u001b[0m in \u001b[0;36mdecode_example\u001b[0;34m(self, value, token_per_repo_id)\u001b[0m\n\u001b[1;32m    191\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    192\u001b[0m             \u001b[0mimage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPIL\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mImage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mBytesIO\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbytes_\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 193\u001b[0;31m         \u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mload\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# to avoid \"Too many open files\" errors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    194\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgetexif\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPIL\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mImage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExifTags\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mBase\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOrientation\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    195\u001b[0m             \u001b[0mimage\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPIL\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mImageOps\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexif_transpose\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.10/site-packages/PIL/ImageFile.py\u001b[0m in \u001b[0;36mload\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    388\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    389\u001b[0m                             \u001b[0mb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mb\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0ms\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 390\u001b[0;31m                             \u001b[0mn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0merr_code\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdecoder\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    391\u001b[0m                             \u001b[0;32mif\u001b[0m \u001b[0mn\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    392\u001b[0m                                 \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for d in join_ds:\n",
    "    print(d[\"instruction\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c87b31e6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f2e2ed5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "\n",
    "save_base_dir = Path(\"/data/regression_output\")\n",
    "save_base_dir.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3502393",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf6debd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "foundation = QwenImageFoundationSaveInterm(QwenConfig())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f55c969",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4abc5c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "\n",
    "for idx, input_data in enumerate(join_ds):\n",
    "\n",
    "    output_dict = foundation.base_pipe(foundation.INPUT_MODEL(\n",
    "        image=[input_data[\"input_img\"]],\n",
    "        prompt=input_data[\"instruction\"],\n",
    "    ))\n",
    "\n",
    "    torch.save(output_dict, save_base_dir/f\"{idx:06d}.pt\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ca95aee",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb8ed336",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dict = torch.load(save_base_dir/f\"{idx:06d}.pt\", weights_only=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69da98cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4329df30",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ind = 10\n",
    "\n",
    "latents_i_start = output_dict[f\"latents_{test_ind}_start\"]\n",
    "t_i = output_dict[f\"t_{test_ind}\"]\n",
    "v_i = output_dict[f\"noise_pred_{test_ind}\"]\n",
    "\n",
    "proj_out_i = latents_i_start - t_i * v_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0996f93b",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_out_i_1d = proj_out_i\n",
    "proj_out_i_2d = foundation.unpack_latents(proj_out_i_1d, output_dict[\"height\"] // 16, output_dict[\"width\"] // 16, )\n",
    "proj_out_i_pil = foundation.latents_to_pil(proj_out_i_2d)\n",
    "proj_out_i_pil[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7928efb",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_1d = output_dict[\"image_latents\"]\n",
    "out_2d = foundation.unpack_latents(out_1d, output_dict[\"height\"] // 16, output_dict[\"width\"] // 16, )\n",
    "out_pil = foundation.latents_to_pil(out_2d)\n",
    "# out_pil[0]\n",
    "# join_ds[idx][\"input_img\"]\n",
    "# join_ds[idx][\"instruction\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dc97e6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_out_i_pil[0].size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ae9d4ef",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d6a4490",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0b01c0d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446ed21d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}