Upload luciddreamer.py with huggingface_hub

luciddreamer.py

@@ -75,7 +75,7 @@ class LucidDreamer:
         self.lama = None
         self.current_model = self.default_model
 
-    def load_model(self, model_name, use_lama=False):
+    def load_model(self, model_name, use_lama=True):
         if model_name is None:
             model_name = self.default_model
         if self.current_model == model_name:
@@ -108,33 +108,36 @@ class LucidDreamer:
         self.current_model = model_name
 
     def rgb(self, prompt, image, negative_prompt='', generator=None, num_inference_steps=50, mask_image=None):
+        image_pil = Image.fromarray(np.round(image * 255.).astype(np.uint8))
+        mask_pil = Image.fromarray(np.round((1 - mask_image) * 255.).astype(np.uint8))
         if self.current_model == self.default_model:
             return self.rgb_model(
                 prompt=prompt,
                 negative_prompt=negative_prompt,
                 generator=generator,
                 num_inference_steps=num_inference_steps,
-                image=image,
-                mask_image=mask_image,
+                image=image_pil,
+                mask_image=mask_pil,
             ).images[0]
 
         kwargs = {
             'negative_prompt': negative_prompt,
             'generator': generator,
-            'strength': 0.8,
+            'strength': 0.9,
             'num_inference_steps': num_inference_steps,
             'height': self.cam.H,
             'width': self.cam.W,
         }
 
-        image_np = np.array(image).astype(float) / 255.0
-        mask_np = np.array(mask_image) / 255.0
-        mask_sum = np.clip((image_np.prod(axis=-1) == 0) + (1 - mask_np), 0, 1)
+        # image_np = np.array(image).astype(float) / 255.0
+        # mask_np = 1.0 - np.array(mask_image) / 255.0
+        image_np = np.round(np.clip(image, 0, 1) * 255.).astype(np.uint8)
+        mask_sum = np.clip((image.prod(axis=-1) == 0) + (1 - mask_image), 0, 1)
         mask_padded = pad_mask(mask_sum, 3)
         masked = image_np * np.logical_not(mask_padded[..., None])
 
         if self.lama is not None:
-            lama_image = Image.fromarray(lama(masked, mask_padded).astype(np.uint8))
+            lama_image = Image.fromarray(self.lama(masked, mask_padded).astype(np.uint8))
         else:
             lama_image = image
 
@@ -309,10 +312,15 @@ class LucidDreamer:
             image_in, mask_in = np.zeros((in_res, in_res, 3), dtype=np.uint8), 255*np.ones((in_res, in_res, 3), dtype=np.uint8)
             image_in[int(in_res/2-h_in/2):int(in_res/2+h_in/2), int(in_res/2-w_in/2):int(in_res/2+w_in/2)] = np.array(rgb_cond)
             mask_in[int(in_res/2-h_in/2):int(in_res/2+h_in/2), int(in_res/2-w_in/2):int(in_res/2+w_in/2)] = 0
+
+            image2 = np.array(Image.fromarray(image_in).resize((self.cam.W, self.cam.H))).astype(float) / 255.0
+            mask2 = np.array(Image.fromarray(mask_in).resize((self.cam.W, self.cam.H))).astype(float) / 255.0
             image_curr = self.rgb(
-                prompt=prompt, image=Image.fromarray(image_in).resize((self.cam.W, self.cam.H)),
+                prompt=prompt,
+                image=image2,
                 negative_prompt=negative_prompt, generator=generator,
-                mask_image=Image.fromarray(mask_in).resize((self.cam.W, self.cam.H)))
+                mask_image=mask2,
+            )
 
         else: # if there is a large gap between height and width, do inpainting
             if w_in > h_in:
@@ -379,9 +387,10 @@ class LucidDreamer:
             border_valid_idx = np.where(mask_hf[round_coord_cam2[1], round_coord_cam2[0]] == 1)[0]  # use valid_idx[border_valid_idx] for world1
 
             image_curr = self.rgb(
-                prompt=prompt, image=Image.fromarray(np.round(image2*255.).astype(np.uint8)),
+                prompt=prompt, image=image2, #Image.fromarray(np.round(image2*255.).astype(np.uint8)),
                 negative_prompt=negative_prompt, generator=generator, num_inference_steps=diff_steps,
-                mask_image=Image.fromarray(np.round((1-mask2[:,:])*255.).astype(np.uint8)))
+                mask_image=mask2, #Image.fromarray(np.round((1-mask2[:,:])*255.).astype(np.uint8))
+            )
             depth_curr = self.d(image_curr)
 
 
@@ -502,6 +511,8 @@ class LucidDreamer:
                                                             pixel_coord_camj[0]/pixel_coord_camj[2]<=W-1, 
                                                             pixel_coord_camj[1]/pixel_coord_camj[2]>=0, 
                                                             pixel_coord_camj[1]/pixel_coord_camj[2]<=H-1)))[0]
+                if len(valid_idxj) == 0:
+                    continue
                 pts_depthsj = pixel_coord_camj[-1:, valid_idxj]
                 pixel_coord_camj = pixel_coord_camj[:2, valid_idxj]/pixel_coord_camj[-1:, valid_idxj]
                 round_coord_camj = np.round(pixel_coord_camj).astype(np.int32)