refactor: Update SDXL-LoRA inference pipeline to load multiple adapter weights

product_diffusion_api/routers/sdxl_text_to_image.py

@@ -82,11 +82,13 @@ def pil_to_s3_json(image: Image.Image,file_name) -> str:
 
 
 @lru_cache(maxsize=1)
-def load_pipeline(model_name, adapter_name):
+def load_pipeline(model_name, adapter_name,adapter_name_2):
     pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(
         "cuda"
     )
     pipe.load_lora_weights(adapter_name)
+    pipe.load_lora_weights(adapter_name_2)
+    pipe.set_adapters([adapter_name, adapter_name_2], adapter_weights=[0.7, 0.8])
     pipe.unload_lora_weights()
     pipe.unet.to(memory_format=torch.channels_last)
     pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
@@ -96,7 +98,7 @@ def load_pipeline(model_name, adapter_name):
     return pipe
 
 
-loaded_pipeline = load_pipeline(config.MODEL_NAME, config.ADAPTER_NAME)
+loaded_pipeline = load_pipeline(config.MODEL_NAME, config.ADAPTER_NAME,config.ADAPTER_NAME_2)
 
 
 # SDXLLoraInference class for running inference

scripts/config.py

@@ -1,9 +1,10 @@
 MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
 ADAPTER_NAME = "VikramSingh178/sdxl-lora-finetune-product-caption"
+ADAPTER_NAME_2 = "VikramSingh178/Products10k-SDXL-Lora"
 VAE_NAME= "madebyollin/sdxl-vae-fp16-fix"
 DATASET_NAME= "hahminlew/kream-product-blip-captions"
 PROJECT_NAME = "Product Photography"
-PRODUCTS_10k_DATASET = "amaye15/Products-10k"
+PRODUCTS_10k_DATASET = "VikramSingh178/Products-10k-BLIP-captions"
 CAPTIONING_MODEL_NAME = "Salesforce/blip-image-captioning-base"
 
 
@@ -14,7 +15,7 @@ class Config:
         self.pretrained_vae_model_name_or_path = VAE_NAME
         self.revision = None
         self.variant = None
-        self.dataset_name = DATASET_NAME
+        self.dataset_name = PRODUCTS_10k_DATASET
         self.dataset_config_name = None
         self.train_data_dir = None
         self.image_column = 'image'

scripts/sdxl_lora_tuner.py

@@ -51,16 +51,8 @@ from diffusers.utils.torch_utils import is_compiled_module
 logger = get_logger(__name__)
 
 
+
 def save_model_card(
-    repo_id: str,
-    images: list = None,
-    base_model: str = None,
-    dataset_name: str = None,
-    train_text_encoder: bool = False,
-    repo_folder: str = None,
-    vae_path: str = None,
-):
-    def save_model_card(
         repo_id: str,
         images: list = None,
         base_model: str = None,
@@ -533,7 +525,7 @@ def main():
     else:
         data_files = {}
         if config.train_data_dir is not None:
-            data_files["train"] = os.path.join(config.train_data_dir, "**")
+            data_files["test"] = os.path.join(config.train_data_dir, "**")
         dataset = load_dataset(
             "imagefolder",
             data_files=data_files,
@@ -544,7 +536,7 @@ def main():
 
     # Preprocessing the datasets.
     # We need to tokenize inputs and targets.
-    column_names = dataset["train"].column_names
+    column_names = dataset["test"].column_names
 
     # 6. Get the column names for input/target.
     DATASET_NAME_MAPPING = {
@@ -651,13 +643,13 @@ def main():
 
     with accelerator.main_process_first():
         if config.max_train_samples is not None:
-            dataset["train"] = (
-                dataset["train"]
+            dataset["test"] = (
+                dataset["test"]
                 .shuffle(seed=config.seed)
                 .select(range(config.max_train_samples))
             )
         # Set the training transforms
-        train_dataset = dataset["train"].with_transform(
+        train_dataset = dataset["test"].with_transform(
             preprocess_train, output_all_columns=True
         )