Update app.py

app.py

@@ -220,7 +220,7 @@ def generate_one_shot(src_img, img_prompt):
 
 def generate_zero_shot(src_img, txt_prompt):
     # init model
-    state_dict = torch.load(f"./checkpoints/{img_prompt[-2:]}/epoch_latest.pth", map_location='cpu')
+    state_dict = torch.load(f"./checkpoints/{txt_prompt.replace(' ', '_')}/epoch_latest.pth", map_location='cpu')
     model = Stylizer(ngf=64, phase=3, model_weights=state_dict['G_ema_model'])
     model.to(device)
     model.eval()
@@ -274,6 +274,7 @@ with gr.Blocks() as demo:
 
         with gr.TabItem("One-Shot"):
             one_shot_src_img = gr.Image(label="Upload Input Face Image", type='filepath', height=400)
+            gr.Examples(examples=["./example/source/01.png", "./example/source/02.png", "./example/source/03.png", "./example/source/04.png"], inputs=one_shot_src_img)
             with gr.Row():
                 gr.Image(shape=(100, 100), value = Image.open("example/reference/01.png"), type='pil', label="ref01")
                 gr.Image(shape=(100, 100), value = Image.open("example/reference/02.png"), type='pil', label="ref02")
@@ -289,6 +290,7 @@ with gr.Blocks() as demo:
         
         with gr.TabItem("Zero-Shot"):
             zero_shot_src_img = gr.Image(label="Upload Input Face Image", type='filepath', height=400)
+            gr.Examples(examples=["./example/source/01.png", "./example/source/02.png", "./example/source/03.png", "./example/source/04.png"], inputs=zero_shot_src_img)
             zero_shot_ref_prompt = gr.Dropdown(
                 label="Txt Prompt",
                 info="Select a reference style prompt",