Update app.py

app.py

@@ -11,11 +11,12 @@ ethnicity_labels = ['African-American', 'American_Indian', 'Black', 'Caucasian',
                     'Southeast_Asian', 'White', 'no_ethnicity_specified']
 models = ['DallE', 'SD_14', 'SD_2']
 nos = [1,2,3,4,5,6,7,8,9,10]
-index = np.load("indexes/knn_10752_65.npy")
+indexes = [768, 1536, 10752]
 ds = load_dataset("tti-bias/identities", split="train")
 
-def get_nearest_64(gender="man", ethnicity="Hispanic", model="SD_14", no=1):
+def get_nearest_64(gender="man", ethnicity="Hispanic", model="SD_14", no=1, index):
     df = ds.remove_columns(["image","image_path"]).to_pandas()
+    index = np.load(f"indexes/knn_{index}_65.npy")
     ix = df.loc[(df['ethnicity'] == ethnicity)  & (df['gender'] == gender) & (df['no'] == no) & (df['model'] == model)].index[0]
     image = ds.select([index[ix][0]])["image"][0]
     neighbors =  ds.select(index[ix][1:25])
@@ -33,6 +34,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             model = gr.Radio(models, label="Model")
+            index = gr.Radio(indexes, label="Visual vocabulary size")
             gender = gr.Radio(gender_labels, label="Gender label")
             no = gr.Radio(nos, label="Image number")
         with gr.Column():
@@ -41,5 +43,5 @@ with gr.Blocks() as demo:
     with gr.Row():
         image = gr.Image()
         gallery = gr.Gallery().style(grid=4)
-    button.click(get_nearest_64, inputs=[gender, ethnicity, model, no], outputs=[image, gallery])
+    button.click(get_nearest_64, inputs=[gender, ethnicity, model, no, index], outputs=[image, gallery])
 demo.launch()