lol maybe

hydra_node/config.py

@@ -175,7 +175,7 @@ def init_config_model():
 
     if config.model_path is None:
         try:
-            config.model_path = f"https://{config.s3_endpoint}/{config.s3_bucket}/{config.s3_file}"
+            config.model_path = f"{config.s3_bucket}/{config.s3_folder}"
             logger.info(f"Path is set to S3 {config.model_path}")
         except:
             logger.error("No model path or S3 info provided")

hydra_node/models.py

@@ -338,19 +338,26 @@ class StableDiffusionModel(nn.Module):
 
     def from_url(self, url):
         #read config url into bytes
-        config_path = self.config.s3_folder + "/config.yaml"
-        s3_file = self.config.s3_folder + "/model.ckpt"
-        model_config = requests.get(config_path, stream='True').raw
-        model_config = OmegaConf.load(model_config)
-        print(f"Downloading model from {url}")
-        
+        s3_file = self.config.model_path + "/config.yaml"
+        headers = web.get_s3_secret_headers(endpoint=self.config.s3_endpoint,
+                                            access_key=self.config.s3_access_key,
+                                            secret_key=self.config.s3_secret_key,
+                                            s3_file=s3_file
+        )
+        url = "https://" + self.config.s3_endpoint + "/" +s3_file
+        tensor_loader = web.CURLStreamFile(url, headers=headers)
+        model_config = OmegaConf.load(tensor_loader)
 
+        s3_file = self.config.model_path + "/model.ckpt"
+        print(f"Downloading model from {url}")
         headers = web.get_s3_secret_headers(endpoint=self.config.s3_endpoint,
                                             access_key=self.config.s3_access_key,
                                             secret_key=self.config.s3_secret_key,
                                             s3_file=s3_file
         )
+        url = "https://" + self.config.s3_endpoint + "/" +s3_file
         tensor_loader = web.CURLStreamFile(url, headers=headers)
+        model = self.load_model_from_config(model_config, tensor_loader)
         return model, model_config
 
     def load_model_from_config(self, config, ckpt, verbose=False):