-
Notifications
You must be signed in to change notification settings - Fork 78
Description
I am encountering that the VRAM usage increases every time i pick a new image and run the generation step it does not release the last cached output and keeps building up until it goes out of memory and starts using CPU&RAM instad .
current setup RTX-3090, 64gb RAM, Ryzen 5900X
example of the issue >
First run snowman : total 14.2GB VRAM used
Second run Dragon : Total 17.2GB VRAM used
Third Run Snowman Again : Total Vram 23.9gb VRAM used
forth Random image: model switches to CPU usage on ram instead 27.5GB used and stays used .
already fixed this localy , not really used to git hub so dont know how to send a change for review so posting the solution i made here.
In app.py:
In the def generate_3d()
exactly before the return function add this should end like this :
del outputs # Explicitly delete output tensors
torch.cuda.empty_cache() # Force release unused memory
return normal_image_display, mesh_path, mesh_path
the lines i added are
del outputs # Explicitly delete output tensors
torch.cuda.empty_cache() # Force release unused memory
here is the whole function for ppl to copy past in :
def generate_3d(image, seed=-1,
ss_guidance_strength=3, ss_sampling_steps=50,
slat_guidance_strength=3, slat_sampling_steps=6,):
if image is None:
return None, None, None
if seed == -1:
seed = np.random.randint(0, MAX_SEED)
image = hi3dgen_pipeline.preprocess_image(image, resolution=1024)
normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
outputs = hi3dgen_pipeline.run(
normal_image,
seed=seed,
formats=["mesh",],
preprocess_image=False,
sparse_structure_sampler_params={
"steps": ss_sampling_steps,
"cfg_strength": ss_guidance_strength,
},
slat_sampler_params={
"steps": slat_sampling_steps,
"cfg_strength": slat_guidance_strength,
},
)
generated_mesh = outputs['mesh'][0]
# Save outputs
import datetime
output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True)
mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb"
# Export mesh
trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
trimesh_mesh.export(mesh_path)
del outputs # Explicitly delete output tensors
torch.cuda.empty_cache() # Force release unused memory
return normal_image, mesh_path, mesh_path