add all

app.py

@@ -58,16 +58,17 @@ def load_models():
 
     # Load custom UNet
     print('Loading custom UNet...')
-    unet_path = "best_21.ckpt"
-    state_dict = torch.load(unet_path, map_location='cpu')
+    # unet_path = "best_21.ckpt"
+    # state_dict = torch.load(unet_path, map_location='cpu')
     
-    # Process the state dict to match the model keys
-    if 'state_dict' in state_dict:
-        new_state_dict = {key.replace('unet.unet.', ''): value for key, value in state_dict['state_dict'].items()}
-        pipeline.unet.load_state_dict(new_state_dict, strict=False)
-    else:
-        pipeline.unet.load_state_dict(state_dict, strict=False)
-        
+    # # Process the state dict to match the model keys
+    # if 'state_dict' in state_dict:
+    #     new_state_dict = {key.replace('unet.unet.', ''): value for key, value in state_dict['state_dict'].items()}
+    #     pipeline.unet.load_state_dict(new_state_dict, strict=False)
+    # else:
+    #     pipeline.unet.load_state_dict(state_dict, strict=False)
+    # pipeline.unet.push_to_hub("YiftachEde/Sharp-It")
+    pipeline.unet = pipeline.unet.from_pretrained("YiftachEde/Sharp-It").to(torch.float16)
     pipeline = pipeline.to(device).to(torch_dtype=torch.float16)
 
     # Load reconstruction model

zero123plus/pipeline.py

@@ -974,7 +974,7 @@ class Zero123PlusPipeline(diffusers.StableDiffusionPipeline):
                     latent_model_input, t
                 )
                 latent_model_input = torch.cat([latent_model_input, cond_latent], dim=1)
-
+                # latent_model_input = latent_model_input.half()
                 # predict the noise residual
                 noise_pred = self.unet(
                     latent_model_input,