Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
 
-dataset = load_dataset("ShoukanLabs/OpenNiji-Dataset")
 
 
 startimg = 0
@@ -9,10 +8,11 @@ startimg = 0
 def get_dataset_forward():
     global startimg
     final = []
-    for idx in range(startimg, startimg + 50):
-        url = dataset["train"]["url"][idx]
-        prompt = dataset["train"]["prompt"][idx]
-        style = dataset["train"]["style"][idx]
+    dataset = load_dataset("ShoukanLabs/OpenNiji-Dataset", split=f"train[{startimg}:{startimg + 50}]")
+    for idx in dataset
+        url = dataset["url"][idx]
+        prompt = dataset["prompt"][idx]
+        style = dataset["style"][idx]
         final.append((url, f"{prompt}\n\n Style: {style}"))
     startimg += 50
     return final