update

app.py

@@ -71,40 +71,42 @@ Best results come from clean, well-lit images with clear subject isolation. Try
 from image_process import prepare_image
 from briarmbg import BriaRMBG
 snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
-rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE)
-rmbg_net.eval()
+#rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE)
+#rmbg_net.eval()
 from triposg.pipelines.pipeline_triposg import TripoSGPipeline
 snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL)
-triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE)
+#triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE)
 
 # mv adapter
 NUM_VIEWS = 6
 from inference_ig2mv_sdxl import prepare_pipeline, preprocess_image, remove_bg
 from mvadapter.utils import get_orthogonal_camera, tensor_to_image, make_image_grid
 from mvadapter.utils.render import NVDiffRastContextWrapper, load_mesh, render
-mv_adapter_pipe = prepare_pipeline(
-    base_model="stabilityai/stable-diffusion-xl-base-1.0",
-    vae_model="madebyollin/sdxl-vae-fp16-fix",
-    unet_model=None,
-    lora_model=None,
-    adapter_path="huanngzh/mv-adapter",
-    scheduler=None,
-    num_views=NUM_VIEWS,
-    device=DEVICE,
-    dtype=torch.float16,
-)
-birefnet = AutoModelForImageSegmentation.from_pretrained(
-        "ZhengPeng7/BiRefNet", trust_remote_code=True
-    )
-birefnet.to(DEVICE)
-transform_image = transforms.Compose(
-    [
-        transforms.Resize((1024, 1024)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ]
-)
-remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
+#mv_adapter_pipe = prepare_pipeline(
+#    base_model="stabilityai/stable-diffusion-xl-base-1.0",
+#    vae_model="madebyollin/sdxl-vae-fp16-fix",
+#    unet_model=None,
+#    lora_model=None,
+#    adapter_path="huanngzh/mv-adapter",
+#    scheduler=None,
+#    num_views=NUM_VIEWS,
+#    device=DEVICE,
+#    dtype=torch.float16,
+#)
+
+
+#birefnet = AutoModelForImageSegmentation.from_pretrained(
+#        "ZhengPeng7/BiRefNet", trust_remote_code=True
+#    )
+#birefnet.to(DEVICE)
+#transform_image = transforms.Compose(
+#    [
+#        transforms.Resize((1024, 1024)),
+#        transforms.ToTensor(),
+#        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+#    ]
+#)
+#remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
 
 if not os.path.exists("checkpoints/RealESRGAN_x2plus.pth"):
     hf_hub_download("dtarnow/UPscaler", filename="RealESRGAN_x2plus.pth", local_dir="checkpoints")
@@ -140,6 +142,8 @@ def run_full(image: str, req: gr.Request):
     
     image_seg = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
 
+    triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE)
+
     outputs = triposg_pipe(
         image=image_seg,
         generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed),
@@ -199,6 +203,19 @@ def run_full(image: str, req: gr.Request):
         .to(DEVICE)
     )
 
+    birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet", trust_remote_code=True
+        )
+    birefnet.to(DEVICE)
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
+
     image = Image.open(image)
     image = remove_bg_fn(image)
     image = preprocess_image(image, height, width)
@@ -207,6 +224,18 @@ def run_full(image: str, req: gr.Request):
     if seed != -1 and isinstance(seed, int):
         pipe_kwargs["generator"] = torch.Generator(device=DEVICE).manual_seed(seed)
 
+    mv_adapter_pipe = prepare_pipeline(
+        base_model="stabilityai/stable-diffusion-xl-base-1.0",
+        vae_model="madebyollin/sdxl-vae-fp16-fix",
+        unet_model=None,
+        lora_model=None,
+        adapter_path="huanngzh/mv-adapter",
+        scheduler=None,
+        num_views=NUM_VIEWS,
+        device=DEVICE,
+        dtype=torch.float16,
+    )
+
     images = mv_adapter_pipe(
         "high quality",
         height=height,
@@ -256,6 +285,9 @@ def run_full(image: str, req: gr.Request):
 @spaces.GPU()
 @torch.no_grad()
 def run_segmentation(image: str):
+    snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL)
+    rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE)
+    rmbg_net.eval()
     image = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
     return image
 
@@ -270,6 +302,7 @@ def image_to_3d(
     target_face_num: int,
     req: gr.Request
 ):
+    triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE)
     outputs = triposg_pipe(
         image=image,
         generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed),
@@ -333,6 +366,19 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
         .to(DEVICE)
     )
 
+    birefnet = AutoModelForImageSegmentation.from_pretrained(
+            "ZhengPeng7/BiRefNet", trust_remote_code=True
+        )
+    birefnet.to(DEVICE)
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize((1024, 1024)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE)
+
     image = Image.open(image)
     image = remove_bg_fn(image)
     image = preprocess_image(image, height, width)
@@ -341,6 +387,18 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
     if seed != -1 and isinstance(seed, int):
         pipe_kwargs["generator"] = torch.Generator(device=DEVICE).manual_seed(seed)
 
+    mv_adapter_pipe = prepare_pipeline(
+        base_model="stabilityai/stable-diffusion-xl-base-1.0",
+        vae_model="madebyollin/sdxl-vae-fp16-fix",
+        unet_model=None,
+        lora_model=None,
+        adapter_path="huanngzh/mv-adapter",
+        scheduler=None,
+        num_views=NUM_VIEWS,
+        device=DEVICE,
+        dtype=torch.float16,
+    )
+
     images = mv_adapter_pipe(
         "high quality",
         height=height,