read model every inference

app.py

@@ -115,9 +115,9 @@ def generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_se
         mesh.apply_transform(rotation_matrix)
         output_path_tmp = os.path.join(output_path, 'source', '{}_{}.obj'.format(prompt_hash, i))
         mesh.export(output_path_tmp)
-        state['latent'].append(latent.clone().detach())
+        state['latent'].append(latent.clone().detach().cpu())
         mesh_path.append(output_path_tmp)
-
+    del latents
     return mesh_path[0], mesh_path[1], mesh_path[2], mesh_path[3], state
 @torch.no_grad()
 def _3d_editing(xm, models, diffusion, initial_noise, start_t, device, instruction, rand_seed, state):
@@ -140,8 +140,18 @@ def _3d_editing(xm, models, diffusion, initial_noise, start_t, device, instructi
     elif 'cyber' in instruction:
         e_type = 'cyber'
 
-    # import pdb; pdb.set_trace()
-    model = models[e_type].to(device)
+    model = load_model('text300M', device=device)
+    with torch.no_grad():
+        new_proj = nn.Linear(1024 * 2, 1024, device=device, dtype=model.wrapped.input_proj.weight.dtype)
+        new_proj.weight = nn.Parameter(torch.zeros_like(new_proj.weight))
+        new_proj.weight[:, :1024].copy_(model.wrapped.input_proj.weight)  #
+        new_proj.bias = nn.Parameter(torch.zeros_like(new_proj.bias))
+        new_proj.bias[:1024].copy_(model.wrapped.input_proj.bias)
+        model.wrapped.input_proj = new_proj
+
+    ckp = torch.load(hf_hub_download(repo_id='silentchen/Shap_Editor', subfolder='single', filename='{}.pt'.format(e_type)), map_location='cpu')
+    model.load_state_dict(ckp['model'])
+
     noise_initial = initial_noise[e_type].to(device)
     noise_start_t = start_t[e_type]
     general_save_path = './logs/edited'
@@ -266,9 +276,9 @@ def main():
     editing_types = ['rainbow', 'santa_hat', 'lego', 'golden', 'wooden', 'cyber']
     # prepare models
     for editing_type in editing_types:
-        tmp_model = load_model('text300M', device=device)
+        tmp_model = load_model('text300M', device=torch.device('cpu'))
         with torch.no_grad():
-            new_proj = nn.Linear(1024 * 2, 1024, device=device, dtype=tmp_model.wrapped.input_proj.weight.dtype)
+            new_proj = nn.Linear(1024 * 2, 1024, device=torch.device('cpu'), dtype=tmp_model.wrapped.input_proj.weight.dtype)
             new_proj.weight = nn.Parameter(torch.zeros_like(new_proj.weight))
             new_proj.weight[:, :1024].copy_(tmp_model.wrapped.input_proj.weight)  #
             new_proj.bias = nn.Parameter(torch.zeros_like(new_proj.bias))
@@ -277,12 +287,13 @@ def main():
 
         ckp = torch.load(hf_hub_download(repo_id='silentchen/Shap_Editor', subfolder='single', filename='{}.pt'.format(editing_type)), map_location='cpu')
         tmp_model.load_state_dict(ckp['model'])
-        noise_initial = ckp['initial_noise']['noise'].to(device)
+        noise_initial = ckp['initial_noise']['noise'].to(torch.device('cpu'))
         initial_noise[editing_type] = noise_initial
         noise_start_t[editing_type] = ckp['t_start']
         models[editing_type] = tmp_model
-    # del models
-    # models = None
+
+    del models
+    models = None
     with Blocks(
             css=css,
             analytics_enabled=False,