update

app.py

@@ -71,7 +71,8 @@ def arbitrary_to_apose(image, seed):
 def apose_to_multiview(apose_img, seed):
     # convert image to PIL.Image
     apose_img = Image.fromarray(apose_img)
-    return infer_api.genStage2(apose_img, seed, num_levels=1)[0]["images"]
+    results, _ = infer_api.genStage2(apose_img, seed, num_levels=1)
+    return results[0]["images"]
 
 def multiview_to_mesh(images):
     mesh_files = infer_api.genStage3(images)
@@ -79,9 +80,9 @@ def multiview_to_mesh(images):
 
 def refine_mesh(apose_img, mesh1, mesh2, mesh3, seed):
     apose_img = Image.fromarray(apose_img)
-    infer_api.genStage2(apose_img, seed, num_levels=2)
-    print(infer_api.multiview_infer.results.keys())
-    refined = infer_api.genStage4([mesh1, mesh2, mesh3], infer_api.multiview_infer.results)
+    _, all_results = infer_api.genStage2(apose_img, seed, num_levels=2)
+    print(all_results.keys())
+    refined = infer_api.genStage4([mesh1, mesh2, mesh3], all_results)
     return refined
 
 with gr.Blocks(title="StdGEN: Semantically Decomposed 3D Character Generation from Single Images") as demo:

infer_api.py

@@ -341,16 +341,6 @@ def run_multiview_infer(data, pipeline, cfg: TestConfig, num_levels=3):
     torch.cuda.empty_cache()    
     return results
 
-@spaces.GPU
-def load_multiview_pipeline(cfg):
-    pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
-        cfg.pretrained_path,
-        torch_dtype=torch.float16,)
-    pipeline.unet.enable_xformers_memory_efficient_attention()
-    if torch.cuda.is_available():
-        pipeline.to(device)
-    return pipeline
-
 
 class InferAPI:
     def __init__(self,
@@ -768,10 +758,13 @@ 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 = {}
+infer_multiview_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
+    infer_multiview_cfg.pretrained_path,
+    torch_dtype=torch.float16,)
+infer_multiview_pipeline.unet.enable_xformers_memory_efficient_attention()
 if torch.cuda.is_available():
     infer_multiview_pipeline.to(device)
+infer_multiview_results = {}
 
 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)),
@@ -791,6 +784,7 @@ def process_im(self, im):
     im = self.image_transforms(im)
     return im
 
+all_results = {}
 
 @spaces.GPU
 def infer_multiview_gen(img, seed, num_levels):
@@ -804,9 +798,9 @@ def infer_multiview_gen(img, seed, num_levels):
     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
+    for k in results:
+        all_results[k] = results[k]
+    return results, all_results
 
 repo_id = "hyz317/StdGEN"
 all_files = list_repo_files(repo_id, revision="main")