update

app.py

@@ -27,6 +27,7 @@ from common import (
     MAX_SEED,
     VERSION,
     active_btn_by_text_content,
+    dispatch_text2image_fn,
     end_session,
     extract_3d_representations_v3,
     extract_urdf,
@@ -35,7 +36,6 @@ from common import (
     get_selected_image,
     image_to_3d,
     start_session,
-    text2image_fn,
 )
 
 app_name = os.getenv("GRADIO_APP")
@@ -395,7 +395,7 @@ with gr.Blocks(delete_cache=(43200, 43200), theme=custom_theme) as demo:
             image_sample3,
         ],
     ).success(
-        text2image_fn,
+        dispatch_text2image_fn,
         inputs=[
             text_prompt,
             img_guidance_scale,

common.py

@@ -28,12 +28,15 @@ _disable_xformers_flash3()
 monkey_path_trellis()
 
 
+import functools
 import gc
 import logging
 import os
 import shutil
 import subprocess
 import sys
+import time
+import traceback
 from glob import glob
 
 import cv2
@@ -91,6 +94,166 @@ logger = logging.getLogger(__name__)
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "false"
 os.environ.setdefault("OPENAI_API_KEY", "sk-placeholder")
 MAX_SEED = 100000
+DIAG_PREFIX = "[T2I-DIAG]"
+
+
+def _diag(event: str, **fields: object) -> None:
+    timestamp = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+    details = " ".join(f"{key}={value!r}" for key, value in fields.items())
+    message = f"{DIAG_PREFIX} {timestamp} event={event}"
+    if details:
+        message = f"{message} {details}"
+    print(message, flush=True)
+
+
+def _diag_cuda(stage: str) -> None:
+    try:
+        cuda_available = torch.cuda.is_available()
+        _diag(
+            "CUDA_STATUS",
+            stage=stage,
+            pid=os.getpid(),
+            cuda_available=cuda_available,
+            cuda_visible_devices=os.getenv("CUDA_VISIBLE_DEVICES"),
+            device_count=torch.cuda.device_count() if cuda_available else 0,
+        )
+        if not cuda_available:
+            return
+
+        device_index = torch.cuda.current_device()
+        free_bytes, total_bytes = torch.cuda.mem_get_info(device_index)
+        properties = torch.cuda.get_device_properties(device_index)
+        _diag(
+            "CUDA_MEMORY",
+            stage=stage,
+            device_index=device_index,
+            device_name=properties.name,
+            capability=f"{properties.major}.{properties.minor}",
+            free_gib=round(free_bytes / 1024**3, 3),
+            total_gib=round(total_bytes / 1024**3, 3),
+            allocated_gib=round(
+                torch.cuda.memory_allocated(device_index) / 1024**3, 3
+            ),
+            reserved_gib=round(
+                torch.cuda.memory_reserved(device_index) / 1024**3, 3
+            ),
+            max_allocated_gib=round(
+                torch.cuda.max_memory_allocated(device_index) / 1024**3, 3
+            ),
+        )
+    except Exception as exc:
+        _diag(
+            "CUDA_STATUS_FAILED",
+            stage=stage,
+            exception_type=type(exc).__name__,
+            exception=str(exc),
+        )
+
+
+def _diag_pipeline(name: str, pipeline: object) -> None:
+    component_details = []
+    components = getattr(pipeline, "components", None)
+    if isinstance(components, dict):
+        for component_name, component in components.items():
+            device = getattr(component, "device", None)
+            dtype = getattr(component, "dtype", None)
+            component_details.append(
+                f"{component_name}:{type(component).__name__}@{device}/{dtype}"
+            )
+
+    _diag(
+        "PIPELINE_STATUS",
+        name=name,
+        pipeline_type=type(pipeline).__name__,
+        device=getattr(pipeline, "device", None),
+        components=";".join(component_details),
+    )
+
+
+def _install_zerogpu_diagnostics() -> None:
+    try:
+        from spaces.zero import client as zero_client
+        from spaces.zero import wrappers as zero_wrappers
+    except Exception as exc:
+        _diag(
+            "ZEROGPU_DIAGNOSTICS_INSTALL_FAILED",
+            exception_type=type(exc).__name__,
+            exception=str(exc),
+        )
+        return
+
+    original_schedule = zero_client.schedule
+    if not getattr(original_schedule, "_embodiedgen_diagnostic", False):
+
+        @functools.wraps(original_schedule)
+        def diagnostic_schedule(*args, **kwargs):
+            started = time.monotonic()
+            _diag(
+                "ZEROGPU_SCHEDULE_ENTER",
+                pid=os.getpid(),
+                task_id=kwargs.get("task_id"),
+                duration=str(kwargs.get("duration")),
+                gpu_size=kwargs.get("gpu_size"),
+            )
+            try:
+                response = original_schedule(*args, **kwargs)
+            except Exception as exc:
+                _diag(
+                    "ZEROGPU_SCHEDULE_EXCEPTION",
+                    exception_type=type(exc).__name__,
+                    exception=str(exc)[:2000],
+                    elapsed_seconds=round(time.monotonic() - started, 3),
+                )
+                raise
+            _diag(
+                "ZEROGPU_SCHEDULE_READY",
+                nvidia_index=getattr(response, "nvidiaIndex", None),
+                idle=getattr(response, "idle", None),
+                elapsed_seconds=round(time.monotonic() - started, 3),
+            )
+            return response
+
+        diagnostic_schedule._embodiedgen_diagnostic = True
+        zero_client.schedule = diagnostic_schedule
+
+    original_worker_init = zero_wrappers.worker_init
+    if not getattr(original_worker_init, "_embodiedgen_diagnostic", False):
+
+        @functools.wraps(original_worker_init)
+        def diagnostic_worker_init(*args, **kwargs):
+            started = time.monotonic()
+            _diag(
+                "ZEROGPU_WORKER_INIT_ENTER",
+                pid=os.getpid(),
+                ppid=os.getppid(),
+                cuda_visible_devices=os.getenv("CUDA_VISIBLE_DEVICES"),
+            )
+            try:
+                result = original_worker_init(*args, **kwargs)
+            except BaseException as exc:
+                _diag(
+                    "ZEROGPU_WORKER_INIT_EXCEPTION",
+                    exception_type=type(exc).__name__,
+                    exception=str(exc)[:2000],
+                    elapsed_seconds=round(time.monotonic() - started, 3),
+                )
+                traceback.print_exc()
+                sys.stderr.flush()
+                raise
+            _diag(
+                "ZEROGPU_WORKER_INIT_READY",
+                result_type=type(result).__name__,
+                elapsed_seconds=round(time.monotonic() - started, 3),
+            )
+            return result
+
+        diagnostic_worker_init._embodiedgen_diagnostic = True
+        zero_wrappers.worker_init = diagnostic_worker_init
+
+    _diag("ZEROGPU_DIAGNOSTICS_INSTALLED", spaces_version="0.51-compatible")
+
+
+_install_zerogpu_diagnostics()
 
 # Global variables for lazy initialization
 _RBG_REMOVER = None
@@ -117,6 +280,8 @@ if os.getenv("GRADIO_APP").startswith("imageto3d"):
     )
     os.makedirs(TMP_DIR, exist_ok=True)
 elif os.getenv("GRADIO_APP").startswith("textto3d"):
+    _pipeline_init_started = time.monotonic()
+    _diag("IMPORT_MODEL_START", name="PIPELINE", model="Sam3dInference")
     if "sam3d" in os.getenv("GRADIO_APP"):
         PIPELINE = Sam3dInference(device="cuda")
     else:
@@ -124,9 +289,39 @@ elif os.getenv("GRADIO_APP").startswith("textto3d"):
             "microsoft/TRELLIS-image-large"
         )
         # PIPELINE.cuda()
+    _diag(
+        "IMPORT_MODEL_READY",
+        name="PIPELINE",
+        model_type=type(PIPELINE).__name__,
+        elapsed_seconds=round(time.monotonic() - _pipeline_init_started, 3),
+    )
     text_model_dir = "weights/Kolors"
+    _pipeline_img_ip_started = time.monotonic()
+    _diag("IMPORT_MODEL_START", name="PIPELINE_IMG_IP", model="Kolors-IP")
     PIPELINE_IMG_IP = build_text2img_ip_pipeline(text_model_dir, ref_scale=0.3)
+    _diag(
+        "IMPORT_MODEL_READY",
+        name="PIPELINE_IMG_IP",
+        model_type=type(PIPELINE_IMG_IP).__name__,
+        elapsed_seconds=round(time.monotonic() - _pipeline_img_ip_started, 3),
+    )
+    _pipeline_img_started = time.monotonic()
+    _diag("IMPORT_MODEL_START", name="PIPELINE_IMG", model="Kolors")
     PIPELINE_IMG = build_text2img_pipeline(text_model_dir)
+    _diag(
+        "IMPORT_MODEL_READY",
+        name="PIPELINE_IMG",
+        model_type=type(PIPELINE_IMG).__name__,
+        elapsed_seconds=round(time.monotonic() - _pipeline_img_started, 3),
+    )
+    _diag(
+        "IMPORT_MODELS_COMPLETE",
+        global_cuda_pipeline_count=3,
+        note=(
+            "Sam3D plus two full Kolors pipelines are resident before "
+            "ZeroGPU dispatch"
+        ),
+    )
     SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
     GEO_CHECKER = MeshGeoChecker(GPT_CLIENT)
     AESTHETIC_CHECKER = ImageAestheticChecker()
@@ -167,9 +362,21 @@ def preprocess_image_fn(
     rmbg_tag: str = "rembg",
     preprocess: bool = True,
 ) -> tuple[Image.Image, Image.Image]:
-    """Preprocess image with lazy model initialization to avoid CUDA init at import time."""
+    """Preprocess an image with lazily initialized background removal."""
     global _RBG_REMOVER, _RBG14_REMOVER
 
+    started = time.monotonic()
+    _diag(
+        "PREPROCESS_ENTER",
+        pid=os.getpid(),
+        rmbg_tag=rmbg_tag,
+        preprocess=preprocess,
+        input_type=type(image).__name__,
+        rembg_cached=_RBG_REMOVER is not None,
+        rmbg14_cached=_RBG14_REMOVER is not None,
+    )
+    _diag_cuda("preprocess_enter")
+
     if isinstance(image, str):
         image = Image.open(image)
     elif isinstance(image, np.ndarray):
@@ -180,18 +387,34 @@ def preprocess_image_fn(
     # Lazy initialization - models are created on first call within @spaces.GPU context
     if rmbg_tag == "rembg":
         if _RBG_REMOVER is None:
+            _diag("PREPROCESS_MODEL_INIT_START", model="RembgRemover")
             _RBG_REMOVER = RembgRemover()
+            _diag("PREPROCESS_MODEL_INIT_READY", model="RembgRemover")
         bg_remover = _RBG_REMOVER
     else:
         if _RBG14_REMOVER is None:
+            _diag("PREPROCESS_MODEL_INIT_START", model="BMGG14Remover")
             _RBG14_REMOVER = BMGG14Remover()
+            _diag("PREPROCESS_MODEL_INIT_READY", model="BMGG14Remover")
         bg_remover = _RBG14_REMOVER
 
+    _diag("PREPROCESS_INFERENCE_START", model=type(bg_remover).__name__)
     image = bg_remover(image)
+    _diag("PREPROCESS_INFERENCE_READY", model=type(bg_remover).__name__)
     image = keep_largest_connected_component(image)
 
     if preprocess:
+        _diag("PREPROCESS_TRELLIS_START")
         image = trellis_preprocess(image)
+        _diag("PREPROCESS_TRELLIS_READY")
+
+    _diag_cuda("preprocess_exit")
+    _diag(
+        "PREPROCESS_EXIT",
+        elapsed_seconds=round(time.monotonic() - started, 3),
+        output_mode=image.mode,
+        output_size=image.size,
+    )
 
     return image, image_cache
 
@@ -581,60 +804,213 @@ def extract_urdf(
     )
 
 
-@spaces.GPU
+@spaces.GPU(duration=180)
 def text2image_fn(
     prompt: str,
     guidance_scale: float,
     infer_step: int = 50,
     ip_image: Image.Image | str = None,
     ip_adapt_scale: float = 0.3,
-    image_wh: int | tuple[int, int] = [1024, 1024],
+    image_wh: int | tuple[int, int] = (1024, 1024),
     rmbg_tag: str = "rembg",
     seed: int = None,
     enable_pre_resize: bool = True,
     n_sample: int = 3,
     req: gr.Request = None,
-):
-    if isinstance(image_wh, int):
-        image_wh = (image_wh, image_wh)
-    output_root = TMP_DIR
-    if req is not None:
-        output_root = os.path.join(output_root, str(req.session_hash))
-        os.makedirs(output_root, exist_ok=True)
-
-    pipeline = PIPELINE_IMG if ip_image is None else PIPELINE_IMG_IP
-    if ip_image is not None:
-        pipeline.set_ip_adapter_scale([ip_adapt_scale])
-
-    images = text2img_gen(
-        prompt=prompt,
-        n_sample=n_sample,
+) -> list[str]:
+    started = time.monotonic()
+    _diag(
+        "WORKER_ENTER",
+        function="text2image_fn",
+        pid=os.getpid(),
+        ppid=os.getppid(),
+        prompt_length=len(prompt) if prompt is not None else None,
         guidance_scale=guidance_scale,
-        pipeline=pipeline,
-        ip_image=ip_image,
-        image_wh=image_wh,
         infer_step=infer_step,
+        has_reference_image=ip_image is not None,
+        ip_adapt_scale=ip_adapt_scale,
+        image_wh=image_wh,
+        rmbg_tag=rmbg_tag,
         seed=seed,
+        enable_pre_resize=enable_pre_resize,
+        n_sample=n_sample,
+        session_hash=getattr(req, "session_hash", None),
+        configured_duration_seconds=180,
     )
+    _diag_cuda("worker_enter")
 
-    for idx in range(len(images)):
-        image = images[idx]
-        images[idx], _ = preprocess_image_fn(
-            image, rmbg_tag, enable_pre_resize
+    try:
+        if isinstance(image_wh, int):
+            image_wh = (image_wh, image_wh)
+        output_root = TMP_DIR
+        if req is not None:
+            output_root = os.path.join(output_root, str(req.session_hash))
+            os.makedirs(output_root, exist_ok=True)
+        _diag("OUTPUT_DIRECTORY_READY", output_root=output_root)
+
+        pipeline_name = "PIPELINE_IMG" if ip_image is None else "PIPELINE_IMG_IP"
+        pipeline = PIPELINE_IMG if ip_image is None else PIPELINE_IMG_IP
+        _diag_pipeline(pipeline_name, pipeline)
+        if ip_image is not None:
+            _diag("IP_ADAPTER_SCALE_START", scale=ip_adapt_scale)
+            pipeline.set_ip_adapter_scale([ip_adapt_scale])
+            _diag("IP_ADAPTER_SCALE_READY", scale=ip_adapt_scale)
+
+        generation_started = time.monotonic()
+        _diag("TEXT2IMAGE_GENERATION_START", pipeline=pipeline_name)
+        images = text2img_gen(
+            prompt=prompt,
+            n_sample=n_sample,
+            guidance_scale=guidance_scale,
+            pipeline=pipeline,
+            ip_image=ip_image,
+            image_wh=image_wh,
+            infer_step=infer_step,
+            seed=seed,
+        )
+        _diag(
+            "TEXT2IMAGE_GENERATION_READY",
+            image_count=len(images),
+            elapsed_seconds=round(time.monotonic() - generation_started, 3),
         )
+        _diag_cuda("after_text2image_generation")
+
+        for idx, image in enumerate(images):
+            preprocess_started = time.monotonic()
+            _diag(
+                "IMAGE_PREPROCESS_START",
+                image_index=idx,
+                image_mode=image.mode,
+                image_size=image.size,
+            )
+            images[idx], _ = preprocess_image_fn(
+                image, rmbg_tag, enable_pre_resize
+            )
+            _diag(
+                "IMAGE_PREPROCESS_READY",
+                image_index=idx,
+                elapsed_seconds=round(
+                    time.monotonic() - preprocess_started, 3
+                ),
+            )
 
-    save_paths = []
-    for idx, image in enumerate(images):
-        save_path = f"{output_root}/sample_{idx}.png"
-        image.save(save_path)
-        save_paths.append(save_path)
+        save_paths = []
+        for idx, image in enumerate(images):
+            save_path = f"{output_root}/sample_{idx}.png"
+            _diag("IMAGE_SAVE_START", image_index=idx, save_path=save_path)
+            image.save(save_path)
+            save_paths.append(save_path)
+            _diag(
+                "IMAGE_SAVE_READY",
+                image_index=idx,
+                save_path=save_path,
+                file_size=os.path.getsize(save_path),
+            )
 
-    logger.info(f"Images saved to {output_root}")
+        _diag(
+            "WORKER_SUCCESS",
+            output_root=output_root,
+            output_count=len(save_paths),
+            elapsed_seconds=round(time.monotonic() - started, 3),
+        )
+        return save_paths + save_paths
+    except Exception as exc:
+        _diag(
+            "WORKER_EXCEPTION",
+            exception_type=type(exc).__name__,
+            exception=str(exc)[:2000],
+            elapsed_seconds=round(time.monotonic() - started, 3),
+        )
+        traceback.print_exc()
+        sys.stderr.flush()
+        _diag_cuda("worker_exception")
+        raise
+    finally:
+        cleanup_started = time.monotonic()
+        _diag("WORKER_CLEANUP_START")
+        try:
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except Exception as exc:
+            _diag(
+                "WORKER_CLEANUP_EXCEPTION",
+                exception_type=type(exc).__name__,
+                exception=str(exc)[:2000],
+            )
+        _diag_cuda("worker_cleanup_complete")
+        _diag(
+            "WORKER_EXIT",
+            total_elapsed_seconds=round(time.monotonic() - started, 3),
+            cleanup_elapsed_seconds=round(
+                time.monotonic() - cleanup_started, 3
+            ),
+        )
 
-    gc.collect()
-    torch.cuda.empty_cache()
 
-    return save_paths + save_paths
+def dispatch_text2image_fn(
+    prompt: str,
+    guidance_scale: float,
+    infer_step: int = 50,
+    ip_image: Image.Image | str = None,
+    ip_adapt_scale: float = 0.3,
+    image_wh: int | tuple[int, int] = (1024, 1024),
+    rmbg_tag: str = "rembg",
+    seed: int = None,
+    enable_pre_resize: bool = True,
+    n_sample: int = 3,
+    req: gr.Request = None,
+) -> list[str]:
+    started = time.monotonic()
+    _diag(
+        "DISPATCH_ENTER",
+        pid=os.getpid(),
+        ppid=os.getppid(),
+        prompt_length=len(prompt) if prompt is not None else None,
+        guidance_scale=guidance_scale,
+        infer_step=infer_step,
+        has_reference_image=ip_image is not None,
+        image_wh=image_wh,
+        rmbg_tag=rmbg_tag,
+        seed=seed,
+        n_sample=n_sample,
+        session_hash=getattr(req, "session_hash", None),
+    )
+    try:
+        result = text2image_fn(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            infer_step=infer_step,
+            ip_image=ip_image,
+            ip_adapt_scale=ip_adapt_scale,
+            image_wh=image_wh,
+            rmbg_tag=rmbg_tag,
+            seed=seed,
+            enable_pre_resize=enable_pre_resize,
+            n_sample=n_sample,
+            req=req,
+        )
+        _diag(
+            "DISPATCH_SUCCESS",
+            output_count=len(result),
+            elapsed_seconds=round(time.monotonic() - started, 3),
+        )
+        return result
+    except Exception as exc:
+        _diag(
+            "DISPATCH_EXCEPTION",
+            exception_type=type(exc).__name__,
+            exception=str(exc)[:2000],
+            elapsed_seconds=round(time.monotonic() - started, 3),
+        )
+        traceback.print_exc()
+        sys.stderr.flush()
+        raise
+    finally:
+        _diag(
+            "DISPATCH_EXIT",
+            elapsed_seconds=round(time.monotonic() - started, 3),
+        )
 
 
 @spaces.GPU