Update app.py

app.py

@@ -12,6 +12,8 @@ import os
 import spaces
 import random
 import torch
+import json
+
 
 from transformers import GLPNFeatureExtractor, GLPNForDepthEstimation
 from transformers import AutoFeatureExtractor, AutoModelForDepthEstimation
@@ -51,6 +53,8 @@ if os.path.exists(counter_file_path):
 else:
     counter = 0
 
+generated_files = []
+
 
 def check_use_custom_or_no(value):
     if value is True:
@@ -236,6 +240,19 @@ def infer(use_custom_model, model_name, weight_name, custom_lora_weight, image_i
     images[0].save(f"viteGradio/images/result{counter}.png")
     print("HELP")
     predict(images[0], counter)
+
+    with open('viteGradio/images/names.json', 'r') as f:
+      filenames = json.load(f)
+
+    result_filename = f"result{counter}.png"
+    depth_filename = f"depth{counter}.png"
+
+    filenames.append(result_filename)
+    filenames.append(depth_filename)
+
+    with open('viteGradio/images/names.json', 'w') as f:
+        json.dump(filenames, f)
+
     counter+=1
     with open(counter_file_path, "w") as file:
       file.write(str(counter))