debug

app.py

@@ -74,11 +74,12 @@ def check_gpu():
     os.environ['LD_LIBRARY_PATH'] = "/usr/local/cuda-12.1/lib64:" + os.environ.get('LD_LIBRARY_PATH', '')
     subprocess.run(['nvidia-smi'])  # Test if CUDA is available
     print(f"torch.cuda.is_available:{torch.cuda.is_available()}")
-    print("Device count:", torch.cuda.device_count())        
+    print("Device count:", torch.cuda.device_count()) 
+
     # test nvdiffrast
     import nvdiffrast.torch as dr
     dr.RasterizeCudaContext(device="cuda:0")
-    print("nvdiffrast initialized successfully")
+    print("nvdiffrast initialized successfully")       
     
 
 # Only check GPU in non-UI debug mode
@@ -163,137 +164,114 @@ def save_py3dmesh_with_trimesh_fast(meshes, save_glb_path=TEMP_MESH_ADDRESS, app
         fix_vert_color_glb(save_glb_path)
     print(f"saving to {save_glb_path}")
 
-# Create model function substitutes for debug mode
-def debug_text_to_detailed(prompt, seed=None):
-    print("DEBUG MODE: text_to_detailed was called")
-    return f"This is a detailed prompt example in debug mode: '{prompt}'"
-
-def debug_text_to_image(prompt, seed=None, strength=1.0,lora_scale=1.0, num_inference_steps=18, redux_hparam=None, init_image=None, **kwargs):
-    print("DEBUG MODE: text_to_image was called")
-    # Return an example image path or create a solid color image
-    example_image = Image.new('RGB', (512, 512), color=(73, 109, 137))
-    return example_image
-
-def debug_bundle_image_to_mesh(*args, **kwargs):
-    print("DEBUG MODE: bundle_image_to_mesh was called")
-    # Return example video and mesh paths
-    return "app_assets/logo_temp_.png", "app_assets/logo_temp_.png"
-
-# Select actual functionality or debug substitutes based on mode
-if UI_ONLY_MODE:
-    text_to_detailed = debug_text_to_detailed
-    text_to_image = debug_text_to_image
-    bundle_image_to_mesh = debug_bundle_image_to_mesh
-    print("UI debug mode functions loaded")
-else:
-    @spaces.GPU
-    def text_to_detailed(prompt, seed=None):
-        # test nvdiffrast
-        import nvdiffrast.torch as dr
-        dr.RasterizeCudaContext(device="cuda:0")
-        print("nvdiffrast initialized successfully")
-
-        print(f"torch.cuda.is_available():{torch.cuda.is_available()}")
-        # print(f"Before text_to_detailed: {torch.cuda.memory_allocated() / 1024**3} GB")
-        return k3d_wrapper.get_detailed_prompt(prompt, seed)
-
-    @spaces.GPU(duration=120)
-    def text_to_image(prompt, seed=None, strength=1.0,lora_scale=1.0, num_inference_steps=18, redux_hparam=None, init_image=None, **kwargs):
-        # subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
-        # print(f"Before text_to_image: {torch.cuda.memory_allocated() / 1024**3} GB")
-        # k3d_wrapper.flux_pipeline.enable_xformers_memory_efficient_attention()
-        k3d_wrapper.renew_uuid()
-        init_image = None
-        # if init_image_path is not None:
-        #     init_image = Image.open(init_image_path)
-        subprocess.run(['nvidia-smi'])  # Test if CUDA is available
-        with torch.no_grad():
-            result = k3d_wrapper.generate_3d_bundle_image_text( 
-                                        prompt,
-                                        image=init_image, 
-                                        strength=strength,
-                                        lora_scale=lora_scale,
-                                        num_inference_steps=num_inference_steps,
-                                        seed=int(seed) if seed is not None else None,
-                                        redux_hparam=redux_hparam,
-                                        save_intermediate_results=True,
-                                        **kwargs)
-        return result[-1]
-
-    @spaces.GPU(duration=120)
-    def image2mesh_preprocess_(input_image_, seed, use_mv_rgb=True):
-        global preprocessed_input_image
-
-        seed = int(seed) if seed is not None else None
-
-        # TODO: delete this later
-        # k3d_wrapper.del_llm_model()
-        
-        input_image_save_path, reference_save_path, caption = image2mesh_preprocess(k3d_wrapper, input_image_, seed, use_mv_rgb)
-
-        preprocessed_input_image = Image.open(input_image_save_path)
-        return reference_save_path, caption
-
-
-    @spaces.GPU(duration=120)
-    def image2mesh_main_(reference_3d_bundle_image, caption, seed, strength1=0.5, strength2=0.95, enable_redux=True, use_controlnet=True, if_video=True):
-        subprocess.run(['nvidia-smi'])  
-        global mesh_cache 
-        seed = int(seed) if seed is not None else None
-
-
-        # TODO: delete this later
-        # k3d_wrapper.del_llm_model()
-
-        input_image = preprocessed_input_image
-
-        reference_3d_bundle_image = torch.tensor(reference_3d_bundle_image).permute(2,0,1)/255
-
-        gen_save_path, recon_mesh_path = image2mesh_main(k3d_wrapper, input_image, reference_3d_bundle_image, caption=caption, seed=seed, strength1=strength1, strength2=strength2, enable_redux=enable_redux, use_controlnet=use_controlnet)
-        mesh_cache = recon_mesh_path
-
-
-        if if_video:
-            video_path = recon_mesh_path.replace('.obj','.mp4').replace('.glb','.mp4')
-            render_video_from_obj(recon_mesh_path, video_path)
-            print(f"After bundle_image_to_mesh: {torch.cuda.memory_allocated() / 1024**3} GB")
-            return gen_save_path, video_path, mesh_cache
-        else:
-            return gen_save_path, recon_mesh_path, mesh_cache
-        # return gen_save_path, recon_mesh_path
-
-    @spaces.GPU(duration=120)
-    def bundle_image_to_mesh(
-            gen_3d_bundle_image, 
-            camera_radius=3.5,
-            lrm_radius = 3.5,
-            isomer_radius = 4.2,
-            reconstruction_stage1_steps = 0,
-            reconstruction_stage2_steps = 50,
-            save_intermediate_results=False
-        ):
-        global mesh_cache
-        print(f"Before bundle_image_to_mesh: {torch.cuda.memory_allocated() / 1024**3} GB")
-        k3d_wrapper.recon_model.init_flexicubes_geometry("cuda:0", fovy=50.0)
-        print(f"init_flexicubes_geometry done")
-        # TODO: delete this later
-        k3d_wrapper.del_llm_model()
-
-        print(f"Before bundle_image_to_mesh after deleting llm model: {torch.cuda.memory_allocated() / 1024**3} GB")
-
-        gen_3d_bundle_image = torch.tensor(gen_3d_bundle_image).permute(2,0,1)/255
-       
-        recon_mesh_path = k3d_wrapper.reconstruct_3d_bundle_image(gen_3d_bundle_image, camera_radius=camera_radius, lrm_render_radius=lrm_radius, isomer_radius=isomer_radius, save_intermediate_results=save_intermediate_results, reconstruction_stage1_steps=int(reconstruction_stage1_steps), reconstruction_stage2_steps=int(reconstruction_stage2_steps))
-        mesh_cache = recon_mesh_path
-        
-        print(f"Mesh generated at: {mesh_cache}")
+@spaces.GPU
+def text_to_detailed(prompt, seed=None):
+    # test nvdiffrast
+    import nvdiffrast.torch as dr
+    dr.RasterizeCudaContext(device="cuda:0")
+    print("nvdiffrast initialized successfully")
+
+    print(f"torch.cuda.is_available():{torch.cuda.is_available()}")
+    # print(f"Before text_to_detailed: {torch.cuda.memory_allocated() / 1024**3} GB")
+    return k3d_wrapper.get_detailed_prompt(prompt, seed)
+
+@spaces.GPU(duration=120)
+def text_to_image(prompt, seed=None, strength=1.0,lora_scale=1.0, num_inference_steps=18, redux_hparam=None, init_image=None, **kwargs):
+    # subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+    # print(f"Before text_to_image: {torch.cuda.memory_allocated() / 1024**3} GB")
+    # k3d_wrapper.flux_pipeline.enable_xformers_memory_efficient_attention()
+    k3d_wrapper.renew_uuid()
+    init_image = None
+    # if init_image_path is not None:
+    #     init_image = Image.open(init_image_path)
+    subprocess.run(['nvidia-smi'])  # Test if CUDA is available
+    with torch.no_grad():
+        result = k3d_wrapper.generate_3d_bundle_image_text( 
+                                    prompt,
+                                    image=init_image, 
+                                    strength=strength,
+                                    lora_scale=lora_scale,
+                                    num_inference_steps=num_inference_steps,
+                                    seed=int(seed) if seed is not None else None,
+                                    redux_hparam=redux_hparam,
+                                    save_intermediate_results=True,
+                                    **kwargs)
+    return result[-1]
+
+@spaces.GPU(duration=120)
+def image2mesh_preprocess_(input_image_, seed, use_mv_rgb=True):
+    global preprocessed_input_image
+
+    seed = int(seed) if seed is not None else None
+
+    # TODO: delete this later
+    # k3d_wrapper.del_llm_model()
+    
+    input_image_save_path, reference_save_path, caption = image2mesh_preprocess(k3d_wrapper, input_image_, seed, use_mv_rgb)
+
+    preprocessed_input_image = Image.open(input_image_save_path)
+    return reference_save_path, caption
+
+
+@spaces.GPU(duration=120)
+def image2mesh_main_(reference_3d_bundle_image, caption, seed, strength1=0.5, strength2=0.95, enable_redux=True, use_controlnet=True, if_video=True):
+    subprocess.run(['nvidia-smi'])  
+    global mesh_cache 
+    seed = int(seed) if seed is not None else None
+
+
+    # TODO: delete this later
+    # k3d_wrapper.del_llm_model()
+
+    input_image = preprocessed_input_image
+
+    reference_3d_bundle_image = torch.tensor(reference_3d_bundle_image).permute(2,0,1)/255
+
+    gen_save_path, recon_mesh_path = image2mesh_main(k3d_wrapper, input_image, reference_3d_bundle_image, caption=caption, seed=seed, strength1=strength1, strength2=strength2, enable_redux=enable_redux, use_controlnet=use_controlnet)
+    mesh_cache = recon_mesh_path
+
+
+    if if_video:
+        video_path = recon_mesh_path.replace('.obj','.mp4').replace('.glb','.mp4')
+        render_video_from_obj(recon_mesh_path, video_path)
+        print(f"After bundle_image_to_mesh: {torch.cuda.memory_allocated() / 1024**3} GB")
+        return gen_save_path, video_path, mesh_cache
+    else:
+        return gen_save_path, recon_mesh_path, mesh_cache
+    # return gen_save_path, recon_mesh_path
+
+@spaces.GPU(duration=120)
+def bundle_image_to_mesh(
+        gen_3d_bundle_image, 
+        camera_radius=3.5,
+        lrm_radius = 3.5,
+        isomer_radius = 4.2,
+        reconstruction_stage1_steps = 0,
+        reconstruction_stage2_steps = 50,
+        save_intermediate_results=False
+    ):
+    global mesh_cache
+    print(f"Before bundle_image_to_mesh: {torch.cuda.memory_allocated() / 1024**3} GB")
+    k3d_wrapper.recon_model.init_flexicubes_geometry("cuda:0", fovy=50.0)
+    print(f"init_flexicubes_geometry done")
+    # TODO: delete this later
+    k3d_wrapper.del_llm_model()
+
+    print(f"Before bundle_image_to_mesh after deleting llm model: {torch.cuda.memory_allocated() / 1024**3} GB")
+
+    gen_3d_bundle_image = torch.tensor(gen_3d_bundle_image).permute(2,0,1)/255
+    
+    recon_mesh_path = k3d_wrapper.reconstruct_3d_bundle_image(gen_3d_bundle_image, camera_radius=camera_radius, lrm_render_radius=lrm_radius, isomer_radius=isomer_radius, save_intermediate_results=save_intermediate_results, reconstruction_stage1_steps=int(reconstruction_stage1_steps), reconstruction_stage2_steps=int(reconstruction_stage2_steps))
+    mesh_cache = recon_mesh_path
+    
+    print(f"Mesh generated at: {mesh_cache}")
+    
+    # Check if file exists
+    if not os.path.exists(mesh_cache):
+        print(f"Warning: Generated mesh file does not exist: {mesh_cache}")
+        return None, mesh_cache
         
-        # Check if file exists
-        if not os.path.exists(mesh_cache):
-            print(f"Warning: Generated mesh file does not exist: {mesh_cache}")
-            return None, mesh_cache
-            
-        return recon_mesh_path, mesh_cache
+    return recon_mesh_path, mesh_cache
 
 # _HEADER_=f"""
 # <img src="{LOGO_PATH}">