Update app.py

app.py

@@ -227,87 +227,89 @@ def marching_cube(b, text, global_info):
     return path
 
 def infer(prompt, samples, steps, scale, seed, global_info):
-    prompt = prompt.replace('/', '')
-    pl.seed_everything(seed)
-    batch_size = samples
-    with torch.no_grad():
-        noise = None
-        c = model.get_learned_conditioning([prompt])
-        unconditional_c = torch.zeros_like(c)
-        sample, _ = sampler.sample(
-            S=steps,
-            batch_size=batch_size,
-            shape=shape,
-            verbose=False,
-            x_T = noise,
-            conditioning = c.repeat(batch_size, 1, 1),
-            unconditional_guidance_scale=scale,
-            unconditional_conditioning=unconditional_c.repeat(batch_size, 1, 1)
-        )
-        decode_res = model.decode_first_stage(sample)
-
-        big_video_list = []
-
-        global_info['decode_res'] = decode_res
-
-        for b in range(batch_size):
-            def render_img(v):
-                rgb_sample, _ = model.first_stage_model.render_triplane_eg3d_decoder(
-                    decode_res[b:b+1], batch_rays_list[v:v+1].to(device), torch.zeros(1, H, H, 3).to(device),
-                )
-                rgb_sample = to8b(rgb_sample.detach().cpu().numpy())[0]
-                rgb_sample = np.stack(
-                    [rgb_sample[..., 2], rgb_sample[..., 1], rgb_sample[..., 0]], -1
-                )
-                rgb_sample = add_text(rgb_sample, str(b))
-                return rgb_sample
-
-            view_num = len(batch_rays_list)
-            video_list = []
-            for v in tqdm.tqdm(range(view_num//8*3, view_num//8*5, 2)):
-                rgb_sample = render_img(v)
-                video_list.append(rgb_sample)
-            big_video_list.append(video_list)
-        # if batch_size == 2:
-        #     cat_video_list = [
-        #         np.concatenate([big_video_list[j][i] for j in range(len(big_video_list))], 1) \
-        #         for i in range(len(big_video_list[0]))
-        #     ]
-        # elif batch_size > 2:
-        #     if batch_size == 3:
-        #         big_video_list.append(
-        #             [np.zeros_like(f) for f in big_video_list[0]]
-        #         )
-        #     cat_video_list = [
-        #         np.concatenate([
-        #             np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1),
-        #             np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1),
-        #         ], 0) \
-        #         for i in range(len(big_video_list[0]))
-        #     ]
-        # else:
-        #     cat_video_list = big_video_list[0]
-
-        for _ in range(4 - batch_size):
-            big_video_list.append(
-                [np.zeros_like(f) + 255 for f in big_video_list[0]]
+    with torch.cuda.device(1):
+        prompt = prompt.replace('/', '')
+        pl.seed_everything(seed)
+        batch_size = samples
+        with torch.no_grad():
+            noise = None
+            c = model.get_learned_conditioning([prompt])
+            unconditional_c = torch.zeros_like(c)
+            sample, _ = sampler.sample(
+                S=steps,
+                batch_size=batch_size,
+                shape=shape,
+                verbose=False,
+                x_T = noise,
+                conditioning = c.repeat(batch_size, 1, 1),
+                unconditional_guidance_scale=scale,
+                unconditional_conditioning=unconditional_c.repeat(batch_size, 1, 1)
             )
-        cat_video_list = [
-            np.concatenate([
-                np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1),
-                np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1),
-            ], 0) \
-            for i in range(len(big_video_list[0]))
-        ]
-
-        path = f"tmp/{prompt.replace(' ', '_')}_{str(datetime.datetime.now()).replace(' ', '_')}.mp4"
-        imageio.mimwrite(path, np.stack(cat_video_list, 0))
+            decode_res = model.decode_first_stage(sample)
+    
+            big_video_list = []
+    
+            global_info['decode_res'] = decode_res
+    
+            for b in range(batch_size):
+                def render_img(v):
+                    rgb_sample, _ = model.first_stage_model.render_triplane_eg3d_decoder(
+                        decode_res[b:b+1], batch_rays_list[v:v+1].to(device), torch.zeros(1, H, H, 3).to(device),
+                    )
+                    rgb_sample = to8b(rgb_sample.detach().cpu().numpy())[0]
+                    rgb_sample = np.stack(
+                        [rgb_sample[..., 2], rgb_sample[..., 1], rgb_sample[..., 0]], -1
+                    )
+                    rgb_sample = add_text(rgb_sample, str(b))
+                    return rgb_sample
+    
+                view_num = len(batch_rays_list)
+                video_list = []
+                for v in tqdm.tqdm(range(view_num//8*3, view_num//8*5, 2)):
+                    rgb_sample = render_img(v)
+                    video_list.append(rgb_sample)
+                big_video_list.append(video_list)
+            # if batch_size == 2:
+            #     cat_video_list = [
+            #         np.concatenate([big_video_list[j][i] for j in range(len(big_video_list))], 1) \
+            #         for i in range(len(big_video_list[0]))
+            #     ]
+            # elif batch_size > 2:
+            #     if batch_size == 3:
+            #         big_video_list.append(
+            #             [np.zeros_like(f) for f in big_video_list[0]]
+            #         )
+            #     cat_video_list = [
+            #         np.concatenate([
+            #             np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1),
+            #             np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1),
+            #         ], 0) \
+            #         for i in range(len(big_video_list[0]))
+            #     ]
+            # else:
+            #     cat_video_list = big_video_list[0]
+    
+            for _ in range(4 - batch_size):
+                big_video_list.append(
+                    [np.zeros_like(f) + 255 for f in big_video_list[0]]
+                )
+            cat_video_list = [
+                np.concatenate([
+                    np.concatenate([big_video_list[0][i], big_video_list[1][i]], 1),
+                    np.concatenate([big_video_list[2][i], big_video_list[3][i]], 1),
+                ], 0) \
+                for i in range(len(big_video_list[0]))
+            ]
+    
+            path = f"tmp/{prompt.replace(' ', '_')}_{str(datetime.datetime.now()).replace(' ', '_')}.mp4"
+            imageio.mimwrite(path, np.stack(cat_video_list, 0))
 
     return global_info, path
 
 def infer_stage2(prompt, selection, seed, global_info, iters):
     prompt = prompt.replace('/', '')
-    mesh_path = marching_cube(int(selection), prompt, global_info)
+    with torch.cuda.device(1):
+        mesh_path = marching_cube(int(selection), prompt, global_info)
     mesh_name = mesh_path.split('/')[-1][:-4]
     # if2_cmd = f"threefiner if2 --mesh {mesh_path} --prompt \"{prompt}\" --outdir tmp --save {mesh_name}_if2.glb --text_dir --front_dir=-y"
     # print(if2_cmd)