Update handler.py

handler.py

@@ -1,4 +1,5 @@
 from typing import  Dict, List, Any
+from PIL import Image
 import torch
 from torch import autocast
 from tqdm.auto import tqdm
@@ -24,13 +25,18 @@ class EndpointHandler():
         # load the optimized model
         print('creating base model...')
         
+        print('creating base model...')
         self.base_name = 'base40M-textvec'
-        #self.base_name = 'base40M'
-        
         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...')
+        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()
@@ -38,6 +44,7 @@ class EndpointHandler():
         
         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))
@@ -58,27 +65,43 @@ class EndpointHandler():
             print('image data found')
         else:
             print('no image data found')
+
         
         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
-        )
+       if use_image:
+           sampler = PointCloudSampler(
+                device=device,
+                models=[base_model, upsampler_model],
+                diffusions=[base_diffusion, upsampler_diffusion],
+                num_points=[1024, 4096 - 1024],
+                aux_channels=['R', 'G', 'B'],
+                guidance_scale=[3.0, 3.0],
+            )
 
-        # Set a test prompt to condition on.
-        # prompt = 'A bluebird mid-flight'
+            # Load an image to condition on.
+            img = Image.open('example_data/cube_stack.jpg')
+       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
-            for x in tqdm(sampler.sample_batch_progressive(batch_size=1, model_kwargs=dict(texts=[inputs]))):
-                samples = x
+            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]