handler.py

from typing import  Dict, List, Any
from PIL import Image
import torch
from torch import autocast
from tqdm.auto import tqdm
from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
from point_e.diffusion.sampler import PointCloudSampler
from point_e.models.download import load_checkpoint
from point_e.models.configs import MODEL_CONFIGS, model_from_config
from point_e.util.plotting import plot_point_cloud
from point_e.util.pc_to_mesh import marching_cubes_mesh
from point_e.util.point_cloud import PointCloud
import json
import base64
import numpy as np
from io import BytesIO
import os


# set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if device.type != 'cuda':
    raise ValueError("need to run on GPU")

class EndpointHandler():
    def __init__(self, path=""):
        self.active = True 
        
        # load the optimized model
        print('creating base model...')
        os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
        
        print('creating base model...')
        self.base_name = 'base40M-textvec'
        self.base_model = model_from_config(MODEL_CONFIGS[self.base_name], device)
        self.base_model.eval()
        self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_name])

        print('creating image model...')
        # default - base40M. use base300M or base1B for better results
        self.base_image_name = 'base40M'
        self.base_image_model = model_from_config(MODEL_CONFIGS[self.base_image_name], device)
        self.base_image_model.eval()
        self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_image_name])

        print('creating upsample model...')
        self.upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)
        self.upsampler_model.eval()
        self.upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])
        
        print('downloading base checkpoint...')
        self.base_model.load_state_dict(load_checkpoint(self.base_name, device))
        self.base_image_model.load_state_dict(load_checkpoint(self.base_image_name, device))
        
        print('downloading upsampler checkpoint...')
        self.upsampler_model.load_state_dict(load_checkpoint('upsample', device))
        
        print('creating SDF model...')
        self.sdf_name = 'sdf'
        self.sdf_model = model_from_config(MODEL_CONFIGS[self.sdf_name], device)
        self.sdf_model.eval()
        
        print('loading SDF model...')
        self.sdf_model.load_state_dict(load_checkpoint(self.sdf_name, device))

    def __call__(self, input_data: Any) -> Any:
        
         # Check if input_data is a string and deserialize if necessary
        if isinstance(input_data, str):
            print("input_data is a string, attempting to deserialize...")
            try:
                input_data = json.loads(input_data)  # Convert JSON string to dictionary
            except json.JSONDecodeError as e:
                print(f"Failed to parse JSON: {e}")
                return None  # Handle the error as appropriate

        command = "null"

        if "command" in input_data:
            command = input_data["command"]

        print(f"the command is: {command}")

        #Assume the user app is still running the old version, and send the data back as it is being expected
        #Currently, the App expects a .ply Mesh to be sent back, and will not have a command input sent with it
        if command == "null":
            temp_pc = self.generate_point_cloud(input_data)
            return self.generate_mesh_from_pc(temp_pc)
        elif command == "generate_pc":
            return self.generate_point_cloud(input_data)  
        elif command == "generate_mesh":
            print("generate_mesh command received...")
            raw_pc = input_data.get("raw_pc")
    
            if raw_pc is None:
                print("raw_pc not found in input_data!")
                return None
    
            # Check if raw_pc is a string and deserialize if necessary
            if isinstance(raw_pc, str):
                print("raw_pc is a string, attempting to deserialize...")
                raw_pc = json.loads(raw_pc)
    
            print("Calling generate_mesh_from_pc...")
            return self.generate_mesh_from_pc(raw_pc)
        elif command == "status":
            return self.check_status()
            
        


    def check_status(self) -> bool:
        return self.active


    def generate_point_cloud(self, data: Any) -> Dict[str, Dict[str, float]]:
        print("generate pc called...")
        use_image = False
        
        #Checks if an image key has been provided, and if so, uses the image data instead of text input
        if "image" in data:
            image_data_encoded = data.pop("image")
            use_image = True
            print('image data found')
        else:
            print('no image data found')

        
        inputs = data.pop("inputs", data)

        if use_image:
            sampler = PointCloudSampler(
                device=device,
                models=[self.base_image_model, self.upsampler_model],
                diffusions=[self.base_diffusion, self.upsampler_diffusion],
                num_points=[1024, 4096 - 1024],
                aux_channels=['R', 'G', 'B'],
                guidance_scale=[3.0, 3.0],
            )
        
            # Load an image to condition on.
            image_data = base64.b64decode(image_data_encoded)
        
            # Convert bytes to PIL Image
            img = Image.open(BytesIO(image_data))
        else:
            sampler = PointCloudSampler(
                device=device,
                models=[self.base_model,self.upsampler_model],
                diffusions=[self.base_diffusion, self.upsampler_diffusion],
                num_points=[1024, 4096 - 1024],
                aux_channels=['R', 'G', 'B'],
                guidance_scale=[3.0, 0.0],
                model_kwargs_key_filter=('texts', ''), # Do not condition the upsampler at all
            ) 
 
        # run inference pipeline
        with autocast(device.type):
            samples = None
            if use_image:
                for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(images=[img]))):
                    samples = x
            else:
                for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[inputs]))):
                    samples = x
            
        #image = self.pipe(inputs, guidance_scale=7.5)["sample"][0]  

        pc = sampler.output_to_point_clouds(samples)[0]
        pc_dict = {}

        data_list = pc.coords.tolist()
        json_string = json.dumps(data_list)
        pc_dict['data'] = json_string
        
        # Convert NumPy arrays to Python lists for serializing
        serializable_channels = {key: value.tolist() for key, value in pc.channels.items()}
        
        # Serialize the dictionary to a JSON-formatted string
        channel_data = json.dumps(serializable_channels)
        pc_dict['channels'] = channel_data
        
        return pc_dict


    def generate_mesh_from_pc(self, pc_data: Any) -> Any:
         # Produce a mesh (with vertex colors)
        print("generate mesh called...")

        #De-serialize both the coords and channel data
        coords_list = json.loads(pc_data['data'])
        channels_dict = json.loads(pc_data['channels'])

        # Reconstruct the PointCloud object
        # Make sure to use .items() for the dictionary to output the key-value pairs together
        point_cloud = PointCloud(
            coords=np.array(coords_list, dtype=np.float32),
            channels={name: np.array(array, dtype=np.float32) for name, array in channels_dict.items()}
        )
        
        mesh = marching_cubes_mesh(
            pc=point_cloud,
            model=self.sdf_model,
            batch_size=4096,
            grid_size=32, # increase to 128 for resolution used in evals
            progress=True,
        )

        # Write the mesh to a PLY file to import into some other program.
        with open('mesh.ply', 'wb') as f:
            mesh.write_ply(f)
            print('Mesh saved to ply file, returning mesh data...')
            #print(mesh)
            
        return mesh