Update app.py

app.py

@@ -82,22 +82,23 @@ from huggingface_hub import hf_hub_download
 
 from utils.tool import NormalTransfer, get_background, get_render_cameras_video, load_mipmap, render_frames
 
-device = "cuda"
+device_0 = "cuda:0"
+device_1 = "cuda:1"
 resolution = 512
 save_dir = "./outputs"
 normal_transfer = NormalTransfer()
-isomer_azimuths = torch.from_numpy(np.array([0, 90, 180, 270])).float().to(device)
-isomer_elevations = torch.from_numpy(np.array([5, 5, 5, 5])).float().to(device)
+isomer_azimuths = torch.from_numpy(np.array([0, 90, 180, 270])).float().to(device_1)
+isomer_elevations = torch.from_numpy(np.array([5, 5, 5, 5])).float().to(device_1)
 isomer_radius = 4.5
-isomer_geo_weights = torch.from_numpy(np.array([1, 0.9, 1, 0.9])).float().to(device)
-isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(device)
+isomer_geo_weights = torch.from_numpy(np.array([1, 0.9, 1, 0.9])).float().to(device_1)
+isomer_color_weights = torch.from_numpy(np.array([1, 0.5, 1, 0.5])).float().to(device_1)
 
 # model initialization and loading
 # flux
 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)
+flux_pipe.to(device=device_0, dtype=torch.bfloat16)
 
 
 # lrm
@@ -109,11 +110,11 @@ 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)
+model = model.to(device_1)
 
 @spaces.GPU
 def lrm_reconstructions(image, input_cameras, save_path=None, name="temp", export_texmap=False, if_save_video=False):
-    images = image.unsqueeze(0).to(device)
+    images = image.unsqueeze(0).to(device_1)
     images = v2.functional.resize(images, 512, interpolation=3, antialias=True).clamp(0, 1)
     # breakpoint()
     with torch.no_grad():
@@ -225,7 +226,7 @@ def reconstruct_3d_model(images, prompt):
     normal_multi_view = images[4:, :3, :, :]
     multi_view_mask = get_background(normal_multi_view)
     rgb_multi_view = rgb_multi_view * rgb_multi_view + (1-multi_view_mask)
-    input_cameras = get_flux_input_cameras(batch_size=1, radius=4.2, fov=30).to(device)
+    input_cameras = get_flux_input_cameras(batch_size=1, radius=4.2, fov=30).to(device_1)
     vertices, faces = lrm_reconstructions(rgb_multi_view, input_cameras, save_path=save_dir_path, name='lrm', export_texmap=False, if_save_video=False)
     # local normal to global normal
 
@@ -235,8 +236,8 @@ def reconstruct_3d_model(images, prompt):
     global_normal = global_normal.permute(0,2,3,1)
     rgb_multi_view = rgb_multi_view.permute(0,2,3,1)
     multi_view_mask = multi_view_mask.permute(0,2,3,1).squeeze(-1)
-    vertices = torch.from_numpy(vertices).to(device)
-    faces = torch.from_numpy(faces).to(device)
+    vertices = torch.from_numpy(vertices).to(device_1)
+    faces = torch.from_numpy(faces).to(device_1)
     vertices = vertices @ rotate_x(np.pi / 2, device=vertices.device)[:3, :3]
     vertices = vertices @ rotate_y(np.pi / 2, device=vertices.device)[:3, :3]
 
@@ -283,7 +284,7 @@ def reconstruct_3d_model(images, prompt):
 @spaces.GPU
 def gradio_pipeline(prompt, seed):
     global model
-    model.init_flexicubes_geometry(device, fovy=50.0)
+    model.init_flexicubes_geometry(device_1, fovy=50.0)
     model = model.eval()
     # 生成多视图图像
     rgb_normal_grid = generate_multi_view_images(prompt, seed)