Update app.py

app.py

@@ -8,11 +8,14 @@ import trimesh
 import random
 from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download, snapshot_download, HfApi
 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')
@@ -25,12 +28,15 @@ 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}")
+logger.info(f"Using device: {DEVICE}, ZeroGPU: {IS_ZEROGPU}")
 
-DEFAULT_FACE_NUMBER = 100000
+DEFAULT_FACE_NUMBER = 50000  # Reduced for L4 and ZeroGPU
 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"
@@ -93,7 +99,7 @@ try:
     ).to(DEVICE)
     transform_image = transforms.Compose(
         [
-            transforms.Resize((1024, 1024)),
+            transforms.Resize((512, 512)),  # Reduced for L4 and ZeroGPU
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ]
@@ -117,10 +123,57 @@ def get_random_hex():
     random_hex = random_bytes.hex()
     return random_hex
 
-@spaces.GPU(duration=3)
-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):
+# 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("HF_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"):
     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,
@@ -141,10 +194,11 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, 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 = 768, 768
+        height, width = 512, 512  # Reduced for L4 and ZeroGPU
         cameras = get_orthogonal_camera(
             elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
             distance=[1.8] * NUM_VIEWS,
@@ -187,11 +241,12 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, 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(
-            "high quality",
+            prompt,
             height=height,
             width=width,
-            num_inference_steps=15,
+            num_inference_steps=10,  # Reduced for L4 and ZeroGPU
             guidance_scale=3.0,
             num_images_per_prompt=NUM_VIEWS,
             control_image=control_images,
@@ -204,6 +259,7 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, 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)
@@ -220,18 +276,19 @@ def run_full(image: str, seed: int = 0, num_inference_steps: int = 50, guidance_
             save_dir=save_dir,
             save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb",
             uv_unwarp=True,
-            uv_size=4096,
+            uv_size=2048,  # Reduced for L4 and ZeroGPU
             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 = 50, guidance_scale: float = 7.5, simplify: bool = True, target_face_num: int = DEFAULT_FACE_NUMBER):
+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"):
     try:
         logger.info("Starting gradio_generate")
         # Verify API key
@@ -255,7 +312,7 @@ def gradio_generate(image: str, seed: int = 0, num_inference_steps: int = 50, 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)
+        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)
         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)}"}
@@ -291,7 +348,6 @@ 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:
@@ -307,7 +363,7 @@ def download_image(url: str, save_path: str) -> str:
         logger.error(f"Failed to download image from {url}: {str(e)}")
         raise
 
-@spaces.GPU()
+@conditional_gpu_decorator()
 @torch.no_grad()
 def run_segmentation(image):
     try:
@@ -332,15 +388,17 @@ 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
 
-@spaces.GPU(duration=3)
+@conditional_gpu_decorator(duration=5)
+@retry_on_gpu_abort(max_attempts=3, delay=5)
 @torch.no_grad()
 def image_to_3d(
-    image,  # Changed to accept FileData dict or PIL Image
+    image,
     seed: int,
     num_inference_steps: int,
     guidance_scale: float,
@@ -350,6 +408,8 @@ 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")
@@ -384,6 +444,7 @@ 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
@@ -391,12 +452,15 @@ def image_to_3d(
         logger.error(f"Error in image_to_3d: {str(e)}")
         raise
 
-@spaces.GPU(duration=3)
+@conditional_gpu_decorator(duration=5)
+@retry_on_gpu_abort(max_attempts=3, delay=5)
 @torch.no_grad()
-def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
+def run_texture(image, mesh_path: str, seed: int, req: gr.Request, style_filter: str = "None"):
     try:
-        logger.info("Running texture generation")
-        height, width = 768, 768
+        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 and ZeroGPU
         cameras = get_orthogonal_camera(
             elevation_deg=[0, 0, 0, 0, 89.99, -89.99],
             distance=[1.8] * NUM_VIEWS,
@@ -431,7 +495,16 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
             .to(DEVICE)
         )
 
-        image = Image.open(image)
+        # Handle both file path and PIL Image
+        if isinstance(image, str):
+            if image.startswith("http"):
+                temp_image_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}.png")
+                image = download_image(image, temp_image_path)
+            image = Image.open(image)
+        elif not isinstance(image, Image.Image):
+            logger.error(f"Invalid image type: {type(image)}")
+            raise ValueError(f"Expected PIL Image or str (path/URL), got {type(image)}")
+
         image = remove_bg_fn(image)
         image = preprocess_image(image, height, width)
 
@@ -439,11 +512,12 @@ 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)
 
+        prompt = f"high quality, {style_filter.lower()}" if style_filter != "None" else "high quality"
         images = mv_adapter_pipe(
-            "high quality",
+            prompt,
             height=height,
             width=width,
-            num_inference_steps=15,
+            num_inference_steps=10,  # Reduced for L4 and ZeroGPU
             guidance_scale=3.0,
             num_images_per_prompt=NUM_VIEWS,
             control_image=control_images,
@@ -456,6 +530,7 @@ def run_texture(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")
@@ -473,7 +548,7 @@ def run_texture(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=4096,
+            uv_size=2048,  # Reduced for L4 and ZeroGPU
             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]],
@@ -485,11 +560,14 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
         logger.error(f"Error in run_texture: {str(e)}")
         raise
 
-@spaces.GPU(duration=3)
+@conditional_gpu_decorator(duration=10)
+@retry_on_gpu_abort(max_attempts=3, delay=5)
 @torch.no_grad()
-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):
+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"):
     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")
@@ -508,10 +586,10 @@ def run_full_api(image, seed: int = 0, num_inference_steps: int = 50, 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)
+        image_seg, mesh_path, textured_glb_path = run_full(image_path, seed, num_inference_steps, guidance_scale, simplify, target_face_num, req, style_filter)
         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)}"}
+        return image_seg, mesh_path, textured_glb_path
     except Exception as e:
         logger.error(f"Error in run_full_api: {str(e)}")
         raise
@@ -524,10 +602,15 @@ try:
         inputs=[
             gr.Image(type="filepath", label="Image"),
             gr.Number(label="Seed", value=0, precision=0),
-            gr.Number(label="Inference Steps", value=50, precision=0),
+            gr.Number(label="Inference Steps", value=30, 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.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",
+            ),
         ],
         outputs="json",
         api_name="/api/generate"
@@ -653,7 +736,7 @@ try:
                                 minimum=8,
                                 maximum=50,
                                 step=1,
-                                value=50,
+                                value=30,  # Reduced for L4 and ZeroGPU
                                 info="Higher steps enhance detail but increase processing time",
                                 elem_classes="gr-slider"
                             )
@@ -689,7 +772,7 @@ try:
                         for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
                     ],
                     fn=run_full,
-                    inputs=[image_prompts],
+                    inputs=[image_prompts, seed, num_inference_steps, guidance_scale, reduce_face, target_face_num, style_filter],
                     outputs=[seg_image, model_output, textured_model_output],
                     cache_examples=True,
                 )
@@ -716,7 +799,7 @@ try:
         ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
         gen_texture_button.click(
             run_texture,
-            inputs=[image_prompts, model_output, seed],
+            inputs=[image_prompts, model_output, seed, style_filter],
             outputs=[textured_model_output]
         )
         demo.load(start_session)