|
1094 | 1094 | " modified_model = CompiledModelDecorator(original_diffusion_model, keep_prob=1)\n",
|
1095 | 1095 | " model.model.diffusion_model = CModelWrapper(modified_model, model.model.diffusion_model.out_channels)\n",
|
1096 | 1096 | " \n",
|
1097 |
| - " dataset = datasets.load_dataset(\"jovianzm/Pexels-400k\", split=\"train\", streaming=True).shuffle(seed=42).take(subset_size)\n", |
| 1097 | + " dataset = datasets.load_dataset(\"google-research-datasets/conceptual_captions\", trust_remote_code=True, split=\"train\", streaming=True).shuffle(seed=42).take(subset_size)\n", |
1098 | 1098 | " \n",
|
1099 | 1099 | " pbar = tqdm(total=subset_size)\n",
|
1100 | 1100 | " channels = model.model.diffusion_model.out_channels\n",
|
1101 | 1101 | " frames = model.temporal_length\n",
|
1102 | 1102 | " h, w = 256 // 8, 256 // 8\n",
|
1103 | 1103 | " noise_shape = [1, channels, frames, h, w]\n",
|
1104 | 1104 | " for batch in dataset:\n",
|
1105 |
| - " prompt = batch[\"title\"]\n", |
1106 |
| - " image_path = batch[\"thumbnail\"]\n", |
| 1105 | + " prompt = batch[\"caption\"]\n", |
| 1106 | + " image_path = batch[\"image_url\"]\n", |
1107 | 1107 | " image = download_image(image_path)\n",
|
1108 | 1108 | " if image is None:\n",
|
1109 | 1109 | " continue\n",
|
|
0 commit comments