Urgent - Output not in HD quality

Hello Team,

Good Day !!!

We have successfully deployed the source code on both the local system and the live server. We have also obtained the output, but it is not in HD quality. We have tried multiple approaches to resolve this issue.

The server specifications are high, as per your instructions, and we are using an NVIDIA A100.

Here, I am sharing our code and server specifications. Kindly review and let us know if any modifications are needed, or provide an updated version of the code from your end.

Note: Since the high-spec server incurs hourly charges on AWS, we request a solution as soon as possible.

run.py

import argparse import logging import os import time

import numpy as np import rembg import torch import xatlas from PIL import Image

from tsr.system import TSR from tsr.utils import remove_background, resize_foreground, save_video, to_gradio_3d_orientation from tsr.bake_texture import bake_texture

from torch.nn.parallel import DistributedDataParallel as DDP

print(torch.cuda.is_available(), 'cuda.is_available()') # Should print True

class Timer: def init(self): self.items = {} self.time_scale = 1000.0 # ms self.time_unit = "ms"

def start(self, name: str) -> None:
    if torch.cuda.is_available():
        print(torch.cuda.device_count(), 'cuda.device_count()')  # Should print the number of GPUs
        print(torch.cuda.get_device_name(0), 'cuda.get_device_name(0)')

        torch.cuda.synchronize()
    self.items[name] = time.time()
    logging.info(f"{name} ...")

def end(self, name: str) -> None:
    if name not in self.items:
        return
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    start_time = self.items.pop(name)
    delta = time.time() - start_time
    t = delta * self.time_scale
    logging.info(f"{name} finished in {t:.2f}{self.time_unit}.")

timer = Timer()

logging.basicConfig( format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO ) parser = argparse.ArgumentParser() parser.add_argument("image", type=str, nargs="+", help="Path to input image(s).") parser.add_argument( "--device", default="cuda", type=str, help="Device to use. If no CUDA-compatible device is found, will fallback to 'cpu'. Default: 'cuda:0'", ) parser.add_argument( "--pretrained-model-name-or-path", default="stabilityai/TripoSR", type=str, help="Path to the pretrained model. Could be either a huggingface model id or a local path. Default: 'stabilityai/TripoSR'", ) parser.add_argument( "--chunk-size", default=8192, type=int, help="Evaluation chunk size for surface extraction and rendering. Smaller chunk size reduces VRAM usage but increases computation time. 0 for no chunking. Default: 8192", ) parser.add_argument( "--mc-resolution", default=512, type=int, help="Marching cubes grid resolution. Default: 256" ) parser.add_argument( "--no-remove-bg", action="store_true", help="If specified, the background will NOT be automatically removed from the input image. Default: false", ) parser.add_argument( "--foreground-ratio", default=0.85, type=float, help="Ratio of the foreground size to the image size. Only used when --no-remove-bg is not specified. Default: 0.85", ) parser.add_argument( "--output-dir", default="output/", type=str, help="Output directory to save the results. Default: 'output/'", ) parser.add_argument( "--model-save-format", default="obj", type=str, choices=["obj", "glb"], help="Format to save the extracted mesh. Default: 'obj'", ) parser.add_argument( "--bake-texture", action="store_true", help="Bake a texture atlas for the extracted mesh, instead of vertex colors", ) parser.add_argument( "--texture-resolution", default=2048, type=int, help="Texture atlas resolution, only useful with --bake-texture. Default: 2048" ) parser.add_argument( "--render", action="store_true", help="If specified, save a NeRF-rendered video. Default: false", ) args = parser.parse_args()

output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True)

Select device

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("✅ Selected device:", device)

timer.start("Initializing model") model = TSR.from_pretrained( args.pretrained_model_name_or_path, config_name="config.yaml", weight_name="model.ckpt", ) model.renderer.set_chunk_size(args.chunk_size)

Move model to device BEFORE wrapping in DataParallel

model.to(device)

Use DataParallel if multiple GPUs are available

if torch.cuda.device_count() > 1: # model = torch.nn.DataParallel(model)

model = DDP(model, device_ids=[rank], output_device=rank)

timer.end("Initializing model")

timer.start("Processing images") images = []

if args.no_remove_bg: rembg_session = None else: rembg_session = rembg.new_session()

for i, image_path in enumerate(args.image): if args.no_remove_bg: image = np.array(Image.open(image_path).convert("RGB")) else: image = remove_background(Image.open(image_path), rembg_session) image = resize_foreground(image, args.foreground_ratio) image = np.array(image).astype(np.float32) / 255.0 image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5 image = Image.fromarray((image * 255.0).astype(np.uint8)) os.makedirs(os.path.join(output_dir, str(i)), exist_ok=True) image.save(os.path.join(output_dir, str(i), f"input.png")) images.append(image) timer.end("Processing images")

for i, image in enumerate(images): logging.info(f"Running image {i + 1}/{len(images)} ...")

timer.start("Running model")
with torch.no_grad():
    scene_codes = model([image], device=device)  # Removed `device=device`
timer.end("Running model")

if args.render:
    timer.start("Rendering")
    render_images = model.render(scene_codes, n_views=30, return_type="pil")
    for ri, render_image in enumerate(render_images[0]):
        render_image.save(os.path.join(output_dir, str(i), f"render_{ri:03d}.png"))
    save_video(
        render_images[0], os.path.join(output_dir, str(i), f"render.mp4"), fps=30
    )
    timer.end("Rendering")

timer.start("Extracting mesh")
meshes = model.module.extract_mesh(scene_codes, not args.bake_texture, resolution=args.mc_resolution)
timer.end("Extracting mesh")

out_mesh_path = os.path.join(output_dir, str(i), f"mesh.{args.model_save_format}")
if args.bake_texture:
    out_texture_path = os.path.join(output_dir, str(i), "texture.png")

    timer.start("Baking texture")
    meshes[0] = meshes[0]
    scene_codes[0] = scene_codes[0].to(device)  # Ensure tensors are moved properly
    with torch.cuda.amp.autocast():
        bake_output = bake_texture(meshes[0], model, scene_codes[0], args.texture_resolution)

    # bake_output = bake_texture(meshes[0], model, scene_codes[0], args.texture_resolution)
    timer.end("Baking texture")

    timer.start("Exporting mesh and texture")
    meshes[0] = to_gradio_3d_orientation(meshes[0])
    xatlas.export(out_mesh_path, meshes[0].vertices[bake_output["vmapping"]], bake_output["indices"], bake_output["uvs"], meshes[0].vertex_normals[bake_output["vmapping"]])
    Image.fromarray((bake_output["colors"] * 255.0).astype(np.uint8)).transpose(Image.FLIP_TOP_BOTTOM).save(out_texture_path)
    timer.end("Exporting mesh and texture")
else:
    timer.start("Exporting mesh")
    meshes[0].export(out_mesh_path)
    timer.end("Exporting mesh")

print("✅ Processing complete!")

run cmd - xvfb-run -a python run.py examples/image-client.png --texture-resolution 8192 --bake-texture --mc-resolution 550 --foreground-ratio 0.8 --output-dir output/