update

app.py

@@ -57,10 +57,10 @@ def load_model(model_name):
     model.eval()
     return model, data_params
 
-model_dict = {
-    model_name: load_model(model_name)
-    for model_name in model_name_mapping
-}
+# model_dict = {
+#     model_name: load_model(model_name)
+#     for model_name in model_name_mapping
+# }
 
 placeholder_image = Image.new("RGBA", (512, 512), (0, 0, 0, 0))
 
@@ -84,10 +84,14 @@ def process_selected_models(uploaded_image, selected_display_names):
 
         if model_name in selected_internal:
             try:
-                model, data_params = model_dict[model_name]
+                model, data_params = load_model(model_name)
                 pixel_values = pil_to_tensor(uploaded_image, **data_params).unsqueeze(0).to(device)
                 output = model(pixel_values)[0]
                 reconstructed_image = tensor_to_pil(output[0].cpu(), **data_params)
+
+                del model
+                torch.cuda.empty_cache()
+
                 result = gr.update(value=reconstructed_image, label=label)
             except Exception as e:
                 print(f"Error in model {model_name}: {e}")