torch.no_grad

app.py

@@ -58,12 +58,14 @@ class Blocks(gr.Blocks):
             config[k] = v
 
         return config
+@torch.no_grad()
 def optimize_all(xm, models, initial_noise, noise_start_t, diffusion, latent_model, device, prompt, instruction, rand_seed):
     state = {}
     out_gen_1, out_gen_2, out_gen_3, out_gen_4, state = generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state)
     edited_1, edited_2, edited_3, edited_4, state = _3d_editing(xm, models, diffusion, initial_noise, noise_start_t, device, instruction, rand_seed, state)
     print(state)
     return out_gen_1, out_gen_2, out_gen_3, out_gen_4, edited_1, edited_2, edited_3, edited_4
+@torch.no_grad()
 def generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state):
     set_seed(rand_seed)
     batch_size = 4
@@ -117,7 +119,7 @@ def generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_se
         mesh_path.append(output_path_tmp)
 
     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):
     set_seed(rand_seed)
     mesh_path = []
@@ -262,7 +264,7 @@ def main():
     initial_noise = dict()
     noise_start_t = dict()
     editing_types = ['rainbow', 'santa_hat', 'lego', 'golden', 'wooden', 'cyber']
-
+    # prepare models
     for editing_type in editing_types:
         tmp_model = load_model('text300M', device=device)
         with torch.no_grad():
@@ -279,7 +281,8 @@ def main():
         initial_noise[editing_type] = noise_initial
         noise_start_t[editing_type] = ckp['t_start']
         models[editing_type] = tmp_model
-
+    # del models
+    # models = None
     with Blocks(
             css=css,
             analytics_enabled=False,
@@ -352,7 +355,7 @@ def main():
         with gr.Column():
             gr.Examples(
                 examples=[
-                    [   "a corgi",
+                    ["a corgi",
                         "Make the color of it look like rainbow",
                         456,
                     ],
@@ -369,7 +372,7 @@ def main():
 
 
     demo.queue(max_size=10, api_open=False)
-    demo.launch(share=True, show_api=False, show_error=True)
+    demo.launch(share=False, show_api=False, show_error=True)
 
 if __name__ == '__main__':
     main()