Update app.py

app.py

@@ -97,7 +97,7 @@ isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(d
 flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(dtype=torch.bfloat16)
 flux_lora_ckpt_path = hf_hub_download(repo_id="LTT/xxx-ckpt", filename="rgb_normal_large.safetensors", repo_type="model")
 flux_pipe.load_lora_weights(flux_lora_ckpt_path)
-
+flux_pipe.to(device=device, dtype=torch.bfloat16)
 
 
 # lrm
@@ -109,7 +109,7 @@ model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt",
 state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
 state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
 model.load_state_dict(state_dict, strict=True)
-
+model = model.to(device)
 
 @spaces.GPU
 def lrm_reconstructions(image, input_cameras, save_path=None, name="temp", export_texmap=False, if_save_video=False):
@@ -280,8 +280,6 @@ def reconstruct_3d_model(images, prompt):
 
 # Gradio 接口函数
 def gradio_pipeline(prompt, seed):
-    flux_pipe.to(device=device, dtype=torch.bfloat16)
-    model = model.to(device)
     model.init_flexicubes_geometry(device, fovy=50.0)
     model = model.eval()
     # 生成多视图图像