Update app.py

app.py

@@ -2,11 +2,17 @@ import os
 import random
 import gradio as gr
 
-def fake_gan():
+def fake_gan_jpg():
     base_dir = os.getcwd()  # Get the current base directory
     jpg_files = [file for file in os.listdir(base_dir) if file.lower().endswith((".jpg", ".jpeg"))]  # List all files ending with ".jpg" or ".jpeg"
     images = [(random.choice(jpg_files), os.path.splitext(file)[0]) for file in jpg_files]
     return images
+    
+def fake_gan_png():
+    base_dir = os.getcwd()  # Get the current base directory
+    files = [file for file in os.listdir(base_dir) if file.lower().endswith((".png"))]  # List all files ending with 
+    images = [(random.choice(files), os.path.splitext(file)[0]) for file in files]
+    return images
 
 with gr.Blocks() as demo:
     with gr.Column(variant="panel"):
@@ -25,7 +31,7 @@ with gr.Blocks() as demo:
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
-    btn.click(fake_gan, None, gallery)
+    btn.click(fake_gan_png, None, gallery)
 
 if __name__ == "__main__":
     demo.launch()