huaweilin commited on
Commit
5a9bf9d
·
1 Parent(s): b460782
Files changed (1) hide show
  1. app.py +9 -5
app.py CHANGED
@@ -57,10 +57,10 @@ def load_model(model_name):
57
  model.eval()
58
  return model, data_params
59
 
60
- model_dict = {
61
- model_name: load_model(model_name)
62
- for model_name in model_name_mapping
63
- }
64
 
65
  placeholder_image = Image.new("RGBA", (512, 512), (0, 0, 0, 0))
66
 
@@ -84,10 +84,14 @@ def process_selected_models(uploaded_image, selected_display_names):
84
 
85
  if model_name in selected_internal:
86
  try:
87
- model, data_params = model_dict[model_name]
88
  pixel_values = pil_to_tensor(uploaded_image, **data_params).unsqueeze(0).to(device)
89
  output = model(pixel_values)[0]
90
  reconstructed_image = tensor_to_pil(output[0].cpu(), **data_params)
 
 
 
 
91
  result = gr.update(value=reconstructed_image, label=label)
92
  except Exception as e:
93
  print(f"Error in model {model_name}: {e}")
 
57
  model.eval()
58
  return model, data_params
59
 
60
+ # model_dict = {
61
+ # model_name: load_model(model_name)
62
+ # for model_name in model_name_mapping
63
+ # }
64
 
65
  placeholder_image = Image.new("RGBA", (512, 512), (0, 0, 0, 0))
66
 
 
84
 
85
  if model_name in selected_internal:
86
  try:
87
+ model, data_params = load_model(model_name)
88
  pixel_values = pil_to_tensor(uploaded_image, **data_params).unsqueeze(0).to(device)
89
  output = model(pixel_values)[0]
90
  reconstructed_image = tensor_to_pil(output[0].cpu(), **data_params)
91
+
92
+ del model
93
+ torch.cuda.empty_cache()
94
+
95
  result = gr.update(value=reconstructed_image, label=label)
96
  except Exception as e:
97
  print(f"Error in model {model_name}: {e}")