Upload handler.py

Created custom Handler

handler.py

@@ -0,0 +1,80 @@
+from typing import  Dict, List, Any
+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
+import base64
+from io import BytesIO
+
+
+# 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=""):
+        # load the optimized model
+        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 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))
+        
+        print('downloading upsampler checkpoint...')
+        self.upsampler_model.load_state_dict(load_checkpoint('upsample', device))
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`dict`:. plotly json Data
+        """
+        inputs = data.pop("inputs", data)
+
+        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
+        )
+
+        # Set a test prompt to condition on.
+        # prompt = 'A bluebird mid-flight'
+ 
+        # run inference pipeline
+        with autocast(device.type):
+            samples = None
+            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]
+        return {"data": pc}
+        #print(pc)
+
+        # encode image as base 64
+        #buffered = BytesIO()
+        #image.save(buffered, format="JPEG")
+        #img_str = base64.b64encode(buffered.getvalue())
+
+        # postprocess the prediction
+        #return {"image": img_str.decode()}