update space

app.py

@@ -153,9 +153,14 @@ 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=torch.device('cpu'))
+        tmp_model = model_from_config(load_config('text300M'), device=device)
+        # print(model_name, kwargs)
+        # print(model)
+
+        # xm = load_model('transmitter', de
+        tmp_model = load_model('text300M', device=device)
         with torch.no_grad():
-            new_proj = nn.Linear(1024 * 2, 1024, device=torch.device('cpu'), dtype=tmp_model.wrapped.input_proj.weight.dtype)
+            new_proj = nn.Linear(1024 * 2, 1024, device=device, 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))
@@ -164,10 +169,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(torch.device('cpu'))
+        tmp_model.eval()
+        # print("loaded latent model")
+        tmp_model.to(device)
+        noise_initial = ckp['initial_noise']['noise'].to(device)
         initial_noise[editing_type] = noise_initial
         noise_start_t[editing_type] = ckp['t_start']
-        models[editing_type] = tmp_model
+        models[editing_type] = tmp_model.to(device)
     @torch.no_grad()
     def optimize_all(prompt, instruction,
                      rand_seed):
@@ -279,12 +287,14 @@ def main():
         os.makedirs(general_save_path, exist_ok=True)
         for i, latent in enumerate(state['latent']):
             latent = latent.to(device)
-            text_embeddings_clip = model.cached_model_kwargs(1, dict(texts=[instruction]))
+            text_embeddings_clip = model.cached_model_kwargs(1, dict(texts=[instruction])).to(device)
             print("shape of latent: ", latent.clone().unsqueeze(0).shape, "instruction: ", instruction)
             ref_latent = latent.clone().unsqueeze(0).to(device)
             t_1 = torch.randint(noise_start_t_e_type, noise_start_t_e_type + 1, (1,), device=device).long()
 
             noise_input = diffusion.q_sample(ref_latent, t_1, noise=noise_initial)
+            print("noise_input:", noise_input.device)
+
             out_1 = diffusion.p_mean_variance(model, noise_input, t_1, clip_denoised=True,
                                               model_kwargs=text_embeddings_clip,
                                               condition_latents=ref_latent)