update

infer_api.py

@@ -758,53 +758,55 @@ class InferSlrmAPI:
 
         return mesh_fpath
 
-class InferMultiviewAPI:
-    def __init__(self, config):
-        parser = argparse.ArgumentParser()
-        parser.add_argument("--seed", type=int, default=42)
-        parser.add_argument("--num_views", type=int, default=6)
-        parser.add_argument("--num_levels", type=int, default=3)
-        parser.add_argument("--pretrained_path", type=str, default='./ckpt/StdGEN-multiview-1024')
-        parser.add_argument("--height", type=int, default=1024)
-        parser.add_argument("--width", type=int, default=576)
-        self.cfg = parser.parse_args()
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.pipeline = load_multiview_pipeline(self.cfg)
-        self.results = {}
-        if torch.cuda.is_available():
-            self.pipeline.to(device)
-
-        self.image_transforms = [transforms.Resize(int(max(self.cfg.height, self.cfg.width))),
-                                 transforms.CenterCrop((self.cfg.height, self.cfg.width)),
-                                 transforms.ToTensor(),
-                                 transforms.Lambda(lambda x: x * 2. - 1),
-                                 ]
-        self.image_transforms = transforms.Compose(self.image_transforms)
-
-        prompt_embeds_path = './multiview/fixed_prompt_embeds_6view'
-        self.normal_text_embeds = torch.load(f'{prompt_embeds_path}/normal_embeds.pt')
-        self.color_text_embeds = torch.load(f'{prompt_embeds_path}/clr_embeds.pt')
-        self.total_views = self.cfg.num_views
-
-
-    def process_im(self, im):
-        im = self.image_transforms(im)
-        return im
-
-    def gen(self, img, seed, num_levels):
-        set_seed(seed)
-        data = {}
-
-        cond_im_rgb = self.process_im(img)
-        cond_im_rgb = torch.stack([cond_im_rgb] * self.total_views, dim=0)
-        data["image_cond_rgb"] = cond_im_rgb[None, ...]
-        data["normal_prompt_embeddings"] = self.normal_text_embeds[None, ...]
-        data["color_prompt_embeddings"] = self.color_text_embeds[None, ...]
-
-        results = run_multiview_infer(data, self.pipeline, self.cfg, num_levels=num_levels)
-        for k in results:
-            self.results[k] = results[k]
-        return results
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--seed", type=int, default=42)
+parser.add_argument("--num_views", type=int, default=6)
+parser.add_argument("--num_levels", type=int, default=3)
+parser.add_argument("--pretrained_path", type=str, default='./ckpt/StdGEN-multiview-1024')
+parser.add_argument("--height", type=int, default=1024)
+parser.add_argument("--width", type=int, default=576)
+infer_multiview_cfg = parser.parse_args()
+infer_multiview_device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+infer_multiview_pipeline = load_multiview_pipeline(infer_multiview_cfg)
+infer_multiview_results = {}
+if torch.cuda.is_available():
+    infer_multiview_pipeline.to(device)
+
+infer_multiview_image_transforms = [transforms.Resize(int(max(infer_multiview_cfg.height, infer_multiview_cfg.width))),
+                            transforms.CenterCrop((infer_multiview_cfg.height, infer_multiview_cfg.width)),
+                            transforms.ToTensor(),
+                            transforms.Lambda(lambda x: x * 2. - 1),
+                            ]
+infer_multiview_image_transforms = transforms.Compose(infer_multiview_image_transforms)
+
+prompt_embeds_path = './multiview/fixed_prompt_embeds_6view'
+infer_multiview_normal_text_embeds = torch.load(f'{prompt_embeds_path}/normal_embeds.pt')
+infer_multiview_color_text_embeds = torch.load(f'{prompt_embeds_path}/clr_embeds.pt')
+infer_multiview_total_views = infer_multiview_cfg.num_views
+
+
+@spaces.GPU
+def process_im(self, im):
+    im = self.image_transforms(im)
+    return im
+
+
+@spaces.GPU
+def infer_multiview_gen(img, seed, num_levels):
+    set_seed(seed)
+    data = {}
+
+    cond_im_rgb = process_im(img)
+    cond_im_rgb = torch.stack([cond_im_rgb] * infer_multiview_total_views, dim=0)
+    data["image_cond_rgb"] = cond_im_rgb[None, ...]
+    data["normal_prompt_embeddings"] = infer_multiview_normal_text_embeds[None, ...]
+    data["color_prompt_embeddings"] = infer_multiview_color_text_embeds[None, ...]
+
+    results = run_multiview_infer(data, infer_multiview_pipeline, infer_multiview_cfg, num_levels=num_levels)
+    # for k in results:
+    #     self.results[k] = results[k]
+    return results
 
 repo_id = "hyz317/StdGEN"
 all_files = list_repo_files(repo_id, revision="main")
@@ -824,7 +826,6 @@ print(f"Using device!!!!!!!!!!!!: {infer_canonicalize_device}", file=sys.stderr)
 infer_canonicalize_config_path = infer_canonicalize_config['config_path']
 infer_canonicalize_loaded_config = OmegaConf.load(infer_canonicalize_config_path)
 
-# infer_canonicalize_setup(**infer_canonicalize_loaded_config)
 
 def infer_canonicalize_setup(
     validation: Dict,