feat: Add Yoso-normal

app.py

@@ -131,7 +131,7 @@ class StableNormal(Geowizard):
     '''
 
     def __init__(self):
-        x_start_pipeline = YOSONormalsPipeline.from_pretrained('Stable-X/yoso-normal-v0-2',  trust_remote_code=True,
+        x_start_pipeline = YOSONormalsPipeline.from_pretrained('Stable-X/yoso-normal-v0-3',  trust_remote_code=True,
                                                                 variant="fp16", torch_dtype=torch.float16)
         self.model = StableNormalPipeline.from_pretrained('Stable-X/stable-normal-v0-1', trust_remote_code=True,
                                                            variant="fp16",  torch_dtype=torch.float16,
@@ -162,6 +162,33 @@ class StableNormal(Geowizard):
 
         return f"model: \n{self.model}"
 
+class YosoNormal(Geowizard):
+    def __init__(self):
+        self.model = YOSONormalsPipeline.from_pretrained('Stable-X/yoso-normal-v0-3',  trust_remote_code=True,
+                                                                variant="fp16", torch_dtype=torch.float16, t_start=0)
+        
+        # two stage concat
+        self.model.x_start_pipeline = x_start_pipeline
+        self.model.x_start_pipeline.to('cuda', torch.float16)
+        self.model.prior.to('cuda', torch.float16)
+
+
+    @torch.no_grad()
+    def __call__(self, img, image_resolution=768):
+        pipe_out = self.model(Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)))
+        pred_normal = pipe_out.prediction[0]
+        pred_normal = (pred_normal + 1) / 2 * 255
+        pred_normal = pred_normal.astype(np.uint8)
+
+        return pred_normal
+
+    def to(self, device):
+        self.model.to(device, torch.float16)
+
+    def __repr__(self):
+
+        return f"model: \n{self.model}"
+
 class DSINE(object):
     '''
     Simple Stable Diffusion Package
@@ -380,8 +407,9 @@ def main():
     marigold_pipe = Marigold()
     geowizard_pipe = Geowizard()
     our_pipe = StableNormal()
+    yoso_pipe = YosoNormal()
 
-    run_demo_server([dsine_pipe, marigold_pipe, geowizard_pipe, our_pipe])
+    run_demo_server([dsine_pipe, marigold_pipe, geowizard_pipe, our_pipe, yoso_pipe])
 
 
 if __name__ == "__main__":

stablenormal/pipeline_yoso_normal.py

@@ -588,7 +588,7 @@ class YOSONormalsPipeline(StableDiffusionControlNetPipeline):
         image_latent = image_latent * self.vae.config.scaling_factor
         image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]
 
-        pred_latent = latents
+        pred_latent = torch.zeros_like(latents)
         if pred_latent is None:
             pred_latent = randn_tensor(
                 image_latent.shape,