Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
-from datasets import load_dataset
+from datasets import load_dataset, DownloadConfig
+
 import random
 import numpy as np
 from transformers import CLIPProcessor, CLIPModel
@@ -46,7 +47,19 @@ clip_model = clip_model.to(device)
 
 # Load the Unsplash dataset
 #dataset = load_dataset("jamescalam/unsplash-25k-photos", split="train", trust_remote_code=True)  # all 25K images are in train split
-dataset = load_dataset("1aurent/unsplash-lite", split="train")
+
+download_config = DownloadConfig(
+    max_retries=5,
+    timeout=60  # Increase from default 10s to 60s
+)
+
+dataset = load_dataset(
+    "1aurent/unsplash-lite", 
+    split="train",
+    download_config=download_config,
+    streaming=True
+)
+
 dataset_size = len(dataset)
 
 # Load gpt and modifed weights for captions