Update app.py

app.py

@@ -8,14 +8,11 @@ import trimesh
 import random
 from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
-from huggingface_hub import hf_hub_download, snapshot_download, HfApi
+from huggingface_hub import hf_hub_download, snapshot_download
 import subprocess
 import shutil
 import base64
 import logging
-import requests
-from functools import wraps
-import time
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -28,15 +25,12 @@ except Exception as e:
     logger.error(f"Failed to install spandrel: {str(e)}")
     raise
 
-# Check if running in ZeroGPU environment
-IS_ZEROGPU = os.getenv("HF_ZERO_SPACE", "0") == "1"
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16
 
-logger.info(f"Using device: {DEVICE}, ZeroGPU: {IS_ZEROGPU}")
+logger.info(f"Using device: {DEVICE}")
 
-DEFAULT_FACE_NUMBER = 50000  # Reduced for L4
+DEFAULT_FACE_NUMBER = 100000
 MAX_SEED = np.iinfo(np.int32).max
 TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git"
 MV_ADAPTER_REPO_URL = "https://github.com/huanngzh/MV-Adapter.git"
@@ -99,7 +93,7 @@ try:
     ).to(DEVICE)
     transform_image = transforms.Compose(
         [
-            transforms.Resize((512, 512)),  # Reduced for L4
+            transforms.Resize((1024, 1024)),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ]
@@ -123,57 +117,10 @@ def get_random_hex():
     random_hex = random_bytes.hex()
     return random_hex
 
-# Retry decorator for GPU tasks
-def retry_on_gpu_abort(max_attempts=3, delay=5):
-    def decorator(func):
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            attempts = 0
-            while attempts < max_attempts:
-                try:
-                    return func(*args, **kwargs)
-                except gr.Error as e:
-                    if "GPU task aborted" in str(e):
-                        attempts += 1
-                        logger.warning(f"GPU task aborted, retrying {attempts}/{max_attempts}")
-                        time.sleep(delay)
-                    else:
-                        raise
-            raise gr.Error("Max retries reached for GPU task")
-        return wrapper
-    return decorator
-
-# Quota check for ZeroGPU
-def check_quota():
-    if not IS_ZEROGPU:
-        return True
-    hf_api = HfApi()
-    try:
-        quota = hf_api.get_space_runtime(token=os.getenv("API_TOKEN"))
-        logger.info(f"Remaining ZeroGPU quota: {quota}")
-        return quota.get("gpu_quota_remaining", 0) > 60
-    except Exception as e:
-        logger.error(f"Failed to check quota: {str(e)}")
-        return False
-
-# Conditional GPU decorator
-def conditional_gpu_decorator(duration=None):
-    def decorator(func):
-        if IS_ZEROGPU:
-            return spaces.GPU(duration=duration)(func) if duration else spaces.GPU()(func)
-        return func
-    return decorator
-
-@conditional_gpu_decorator(duration=10)
-@retry_on_gpu_abort(max_attempts=3, delay=5)
-def run_full(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER, req=None, style_filter: str = "None"):
+@spaces.GPU(duration=5)
+def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER, req=None):
     try:
-        logger.info(f"Starting run_full with image: {image}, seed: {seed}, style: {style_filter}")
-        if not check_quota():
-            raise gr.Error("Insufficient GPU quota remaining")
         image_seg = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
-        logger.info("Image segmentation completed")
-        logger.info(f"VRAM usage after segmentation: {torch.cuda.memory_allocated(DEVICE)/1e9:.2f} GB")
 
         outputs = triposg_pipe(
             image=image_seg,
@@ -194,11 +141,10 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_
         mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
         mesh.export(mesh_path)
         logger.info(f"Saved mesh to {mesh_path}")
-        logger.info(f"VRAM usage after mesh generation: {torch.cuda.memory_allocated(DEVICE)/1e9:.2f} GB")
 
         torch.cuda.empty_cache()
 
-        height, width = 512, 512  # Reduced for L4
+        height, width = 768, 768
         cameras = get_orthogonal_camera(
             elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
             distance=[1.8] * NUM_VIEWS,
@@ -241,12 +187,11 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_
         if seed != -1 and isinstance(seed, int):
             pipe_kwargs["generator"] = torch.Generator(device=DEVICE).manual_seed(seed)
 
-        prompt = f"high quality, {style_filter.lower()}" if style_filter != "None" else "high quality"
         images = mv_adapter_pipe(
-            prompt,
+            "high quality",
             height=height,
             width=width,
-            num_inference_steps=10,  # Reduced for L4
+            num_inference_steps=15,
             guidance_scale=3.0,
             num_images_per_prompt=NUM_VIEWS,
             control_image=control_images,
@@ -259,7 +204,6 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_
         ).images
 
         torch.cuda.empty_cache()
-        logger.info(f"VRAM usage after texture generation: {torch.cuda.memory_allocated(DEVICE)/1e9:.2f} GB")
         os.makedirs(save_dir, exist_ok=True)
         mv_image_path = os.path.join(save_dir, f"mv_adapter_{get_random_hex()}.png")
         make_image_grid(images, rows=1).save(mv_image_path)
@@ -276,19 +220,18 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_
             save_dir=save_dir,
             save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
             uv_unwarp=True,
-            uv_size=2048,  # Reduced for L4
+            uv_size=4096,
             rgb_path=mv_image_path,
             rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
             camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
         )
 
-        logger.info(f"run_full completed successfully, textured model saved to {textured_glb_path}")
         return image_seg, mesh_path, textured_glb_path
     except Exception as e:
         logger.error(f"Error in run_full: {str(e)}")
         raise
 
-def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 30, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER, style_filter: str = "None"):
+def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 50, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER):
     try:
         logger.info("Starting gradio_generate")
         # Verify API key
@@ -312,7 +255,7 @@ def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 30, gu
                 logger.error(f"Image file not found: {temp_image_path}")
                 raise ValueError("Invalid or missing image file")
 
-        image_seg, mesh_path, textured_glb_path = run_full(temp_image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req=None, style_filter=style_filter)
+        image_seg, mesh_path, textured_glb_path = run_full(temp_image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req=None)
         session_hash = os.path.basename(os.path.dirname(textured_glb_path))
         logger.info(f"Generated model at /files/{session_hash}/{os.path.basename(textured_glb_path)}")
         return {"file_url": f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"}
@@ -348,6 +291,7 @@ def get_random_seed(randomize_seed, seed):
         logger.error(f"Error in get_random_seed: {str(e)}")
         raise
 
+
 def download_image(url: str, save_path: str) -> str:
     """Download an image from a URL and save it locally."""
     try:
@@ -363,7 +307,7 @@ def download_image(url: str, save_path: str) -> str:
         logger.error(f"Failed to download image from {url}: {str(e)}")
         raise
 
-@conditional_gpu_decorator()
+@spaces.GPU()
 @torch.no_grad()
 def run_segmentation(image):
     try:
@@ -388,17 +332,15 @@ def run_segmentation(image):
 
         image = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
         logger.info("Segmentation complete")
-        torch.cuda.empty_cache()
         return image
     except Exception as e:
         logger.error(f"Error in run_segmentation: {str(e)}")
         raise
 
-@conditional_gpu_decorator(duration=5)
-@retry_on_gpu_abort(max_attempts=3, delay=5)
+@spaces.GPU(duration=5)
 @torch.no_grad()
 def image_to_3d(
-    image,
+    image,  # Changed to accept FileData dict or PIL Image
     seed: int,
     num_inference_steps: int,
     guidance_scale: float,
@@ -408,8 +350,6 @@ def image_to_3d(
 ):
     try:
         logger.info("Running image_to_3d")
-        if not check_quota():
-            raise gr.Error("Insufficient GPU quota remaining")
         # Handle FileData dict from gradio_client
         if isinstance(image, dict):
             image_path = image.get("path") or image.get("url")
@@ -444,7 +384,6 @@ def image_to_3d(
         mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb")
         mesh.export(mesh_path)
         logger.info(f"Saved mesh to {mesh_path}")
-        logger.info(f"VRAM usage after mesh generation: {torch.cuda.memory_allocated(DEVICE)/1e9:.2f} GB")
 
         torch.cuda.empty_cache()
         return mesh_path
@@ -452,15 +391,12 @@ def image_to_3d(
         logger.error(f"Error in image_to_3d: {str(e)}")
         raise
 
-@conditional_gpu_decorator(duration=5)
-@retry_on_gpu_abort(max_attempts=3, delay=5)
+@spaces.GPU(duration=5)
 @torch.no_grad()
-def run_texture(image: Image.Image, mesh_path: str, seed: int, req: gr.Request, style_filter: str = "None"):
+def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
     try:
-        logger.info(f"Running texture generation with style: {style_filter}")
-        if not check_quota():
-            raise gr.Error("Insufficient GPU quota remaining")
-        height, width = 512, 512  # Reduced for L4
+        logger.info("Running texture generation")
+        height, width = 768, 768
         cameras = get_orthogonal_camera(
             elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
             distance=[1.8] * NUM_VIEWS,
@@ -495,6 +431,7 @@ def run_texture(image: Image.Image, mesh_path: str, seed: int, req: gr.Request,
             .to(DEVICE)
         )
 
+        image = Image.open(image)
         image = remove_bg_fn(image)
         image = preprocess_image(image, height, width)
 
@@ -502,12 +439,11 @@ def run_texture(image: 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)
 
-        prompt = f"high quality, {style_filter.lower()}" if style_filter != "None" else "high quality"
         images = mv_adapter_pipe(
-            prompt,
+            "high quality",
             height=height,
             width=width,
-            num_inference_steps=10,  # Reduced for L4
+            num_inference_steps=15,
             guidance_scale=3.0,
             num_images_per_prompt=NUM_VIEWS,
             control_image=control_images,
@@ -520,7 +456,6 @@ def run_texture(image: Image.Image, mesh_path: str, seed: int, req: gr.Request,
         ).images
 
         torch.cuda.empty_cache()
-        logger.info(f"VRAM usage after texture generation: {torch.cuda.memory_allocated(DEVICE)/1e9:.2f} GB")
         save_dir = os.path.join(TMP_DIR, str(req.session_hash))
         os.makedirs(save_dir, exist_ok=True)
         mv_image_path = os.path.join(save_dir, f"mv_adapter_{get_random_hex()}.png")
@@ -538,7 +473,7 @@ def run_texture(image: Image.Image, mesh_path: str, seed: int, req: gr.Request,
             save_dir=save_dir,
             save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
             uv_unwarp=True,
-            uv_size=2048,  # Reduced for L4
+            uv_size=4096,
             rgb_path=mv_image_path,
             rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"),
             camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]],
@@ -550,14 +485,11 @@ def run_texture(image: Image.Image, mesh_path: str, seed: int, req: gr.Request,
         logger.error(f"Error in run_texture: {str(e)}")
         raise
 
-@conditional_gpu_decorator(duration=10)
-@retry_on_gpu_abort(max_attempts=3, delay=5)
+@spaces.GPU(duration=5)
 @torch.no_grad()
-def run_full_api(image, seed: int = 0, num_inference_steps: int = 30, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER, req: gr.Request = None, style_filter: str = "None"):
+def run_full_api(image, seed: int = 0, num_inference_steps: int = 50, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER, req: gr.Request = None):
     try:
         logger.info("Running run_full_api")
-        if not check_quota():
-            raise gr.Error("Insufficient GPU quota remaining")
         # Handle FileData dict or URL
         if isinstance(image, dict):
             image_path = image.get("path") or image.get("url")
@@ -576,9 +508,10 @@ def run_full_api(image, seed: int = 0, num_inference_steps: int = 30, guidance_s
                 logger.error(f"Invalid image path: {image_path}")
                 raise ValueError(f"Invalid image path: {image_path}")
 
-        image_seg, mesh_path, textured_glb_path = run_full(image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req, style_filter)
-        logger.info(f"Generated textured model at /files/{os.path.basename(os.path.dirname(textured_glb_path))}/{os.path.basename(textured_glb_path)}")
-        return image_seg, mesh_path, textured_glb_path
+        image_seg, mesh_path, textured_glb_path = run_full(image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req)
+        session_hash = os.path.basename(os.path.dirname(textured_glb_path))
+        logger.info(f"Generated textured model at /files/{session_hash}/{os.path.basename(textured_glb_path)}")
+        return {"file_url": f"/files/{session_hash}/{os.path.basename(textured_glb_path)}"}
     except Exception as e:
         logger.error(f"Error in run_full_api: {str(e)}")
         raise
@@ -591,15 +524,10 @@ try:
         inputs=[
             gr.Image(type="filepath", label="Image"),
             gr.Number(label="Seed", value=0, precision=0),
-            gr.Number(label="Inference Steps", value=30, precision=0),
+            gr.Number(label="Inference Steps", value=50, precision=0),
             gr.Number(label="Guidance Scale", value=7.5),
             gr.Checkbox(label="Simplify Mesh", value=True),
-            gr.Number(label="Target Face Number", value=DEFAULT_FACE_NUMBER, precision=0),
-            gr.Dropdown(
-                choices=["None", "Realistic", "Fantasy", "Cartoon", "Sci-Fi", "Vintage", "Cosmic", "Neon"],
-                label="Style Filter",
-                value="None",
-            ),
+            gr.Number(label="Target Face Number", value=DEFAULT_FACE_NUMBER, precision=0)
         ],
         outputs="json",
         api_name="/api/generate"
@@ -725,7 +653,7 @@ try:
                                 minimum=8,
                                 maximum=50,
                                 step=1,
-                                value=30,  # Reduced for L4
+                                value=50,
                                 info="Higher steps enhance detail but increase processing time",
                                 elem_classes="gr-slider"
                             )
@@ -761,7 +689,7 @@ try:
                         for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
                     ],
                     fn=run_full,
-                    inputs=[image_prompts, seed, num_inference_steps, guidance_scale, reduce_face, target_face_num, style_filter],
+                    inputs=[image_prompts],
                     outputs=[seg_image, model_output, textured_model_output],
                     cache_examples=True,
                 )
@@ -788,7 +716,7 @@ try:
         ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
         gen_texture_button.click(
             run_texture,
-            inputs=[image_prompts, model_output, seed, style_filter],
+            inputs=[image_prompts, model_output, seed],
             outputs=[textured_model_output]
         )
         demo.load(start_session)