Update app.py

app.py

@@ -94,12 +94,11 @@ isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(d
 
 # model initialization and loading
 # flux
-flux_pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=access_token).to(device=device, dtype=torch.bfloat16)
+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)
-generator = torch.Generator(device=device).manual_seed(10)
+
 
 # lrm
 config = OmegaConf.load("./models/lrm/config/PRM_inference.yaml")
@@ -111,9 +110,6 @@ 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)
-model.init_flexicubes_geometry(device, fovy=50.0)
-model = model.eval()
 
 @spaces.GPU
 def lrm_reconstructions(image, input_cameras, save_path=None, name="temp", export_texmap=False, if_save_video=False):
@@ -284,6 +280,10 @@ 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()
     # 生成多视图图像
     rgb_normal_grid = generate_multi_view_images(prompt, seed)
     image_preview = Image.fromarray((rgb_normal_grid * 255).astype(np.uint8))