minor changes

app.py

@@ -30,14 +30,13 @@ EXAMPLE_TEXT_PROMPTS = [
 ]
 
 
-def load_adapted_unet(unet_pretraining_type, exp_path, pipe):
+def load_adapted_unet(unet_pretraining_type, pipe):
 
     """
     Loads the adapted U-Net for the selected PEFT Type
 
     Parameters:
         unet_pretraining_type (str): The type of PEFT to use for generating the X-ray
-        exp_path (str): The path to the best trained model for the selected PEFT Type
         pipe (StableDiffusionPipeline): The Stable Diffusion Pipeline to use for generating the X-ray
 
     Returns:
@@ -45,6 +44,7 @@ def load_adapted_unet(unet_pretraining_type, exp_path, pipe):
     """
 
     sd_folder_path = "runwayml/stable-diffusion-v1-5"
+    exp_path = ''
 
     if unet_pretraining_type == "freeze":
         pass
@@ -67,19 +67,19 @@ def load_adapted_unet(unet_pretraining_type, exp_path, pipe):
         exp_path = os.path.join(exp_path, "pytorch_lora_weights.safetensors")
         pipe.unet.load_attn_procs(exp_path)
     else:
-        exp_path = unet_pretraining_type + "_" + "diffusion_pytorch_model.safetensors"
-        state_dict = load_file(exp_path)
+        # exp_path = unet_pretraining_type + "_" + "diffusion_pytorch_model.safetensors"
+        # state_dict = load_file(exp_path)
+        state_dict = load_file(unet_pretraining_type + "_" + "diffusion_pytorch_model.safetensors")
         print(pipe.unet.load_state_dict(state_dict, strict=False))
 
 
-def loadSDModel(unet_pretraining_type, exp_path, cuda_device):
+def loadSDModel(unet_pretraining_type, cuda_device):
 
     """
     Loads the Stable Diffusion Model for the selected PEFT Type
 
     Parameters:
         unet_pretraining_type (str): The type of PEFT to use for generating the X-ray
-        exp_path (str): The path to the best trained model for the selected PEFT Type
         cuda_device (str): The CUDA device to use for generating the X-ray
 
     Returns:
@@ -90,104 +90,104 @@ def loadSDModel(unet_pretraining_type, exp_path, cuda_device):
 
     pipe = StableDiffusionPipeline.from_pretrained(sd_folder_path, revision="fp16")
 
-    load_adapted_unet(unet_pretraining_type, exp_path, pipe)
+    load_adapted_unet(unet_pretraining_type, pipe)
     pipe.safety_checker = None
 
     return pipe
 
 
-def load_all_pipelines():
-
-    """
-    Loads all the Stable Diffusion Pipelines for each PEFT Type for efficient caching (Design Choice 2)
-
-    Parameters:
-        None
-
-    Returns:
-        sd_pipeline_full (StableDiffusionPipeline): The Stable Diffusion Pipeline for Full Fine-Tuning
-        sd_pipeline_norm (StableDiffusionPipeline): The Stable Diffusion Pipeline for Norm Fine-Tuning
-        sd_pipeline_bias (StableDiffusionPipeline): The Stable Diffusion Pipeline for Bias Fine-Tuning
-        sd_pipeline_attention (StableDiffusionPipeline): The Stable Diffusion Pipeline for Attention Fine-Tuning
-        sd_pipeline_NBA (StableDiffusionPipeline): The Stable Diffusion Pipeline for NBA Fine-Tuning
-        sd_pipeline_difffit (StableDiffusionPipeline): The Stable Diffusion Pipeline for Difffit Fine-Tuning
-    """
-
-    # Dictionary containing the path to the best trained models for each PEFT type
-    MODEL_PATH_DICT = {
-        "full": "full_diffusion_pytorch_model.safetensors",
-        "norm": "norm_diffusion_pytorch_model.safetensors",
-        "bias": "bias_diffusion_pytorch_model.safetensors",
-        "attention": "attention_diffusion_pytorch_model.safetensors",
-        "norm_bias_attention": "norm_bias_attention_diffusion_pytorch_model.safetensors",
-        "difffit": "difffit_diffusion_pytorch_model.safetensors",
-    }
-
-    device = "0"
-    cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu"
-
-    # Full FT
-    unet_pretraining_type = "full"
-    print("Loading Pipeline for Full Fine-Tuning")
-    sd_pipeline_full = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    # Norm
-    unet_pretraining_type = "norm"
-    print("Loading Pipeline for Norm Fine-Tuning")
-    sd_pipeline_norm = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    # bias
-    unet_pretraining_type = "bias"
-    print("Loading Pipeline for Bias Fine-Tuning")
-    sd_pipeline_bias = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    # attention
-    unet_pretraining_type = "attention"
-    print("Loading Pipeline for Attention Fine-Tuning")
-    sd_pipeline_attention = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    # NBA
-    unet_pretraining_type = "norm_bias_attention"
-    print("Loading Pipeline for NBA Fine-Tuning")
-    sd_pipeline_NBA = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    # difffit
-    unet_pretraining_type = "difffit"
-    print("Loading Pipeline for Difffit Fine-Tuning")
-    sd_pipeline_difffit = loadSDModel(
-        unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
-        cuda_device=cuda_device,
-    )
-
-    return (
-        sd_pipeline_full,
-        sd_pipeline_norm,
-        sd_pipeline_bias,
-        sd_pipeline_attention,
-        sd_pipeline_NBA,
-        sd_pipeline_difffit,
-    )
+# def load_all_pipelines():
+
+#     """
+#     Loads all the Stable Diffusion Pipelines for each PEFT Type for efficient caching (Design Choice 2)
+
+#     Parameters:
+#         None
+
+#     Returns:
+#         sd_pipeline_full (StableDiffusionPipeline): The Stable Diffusion Pipeline for Full Fine-Tuning
+#         sd_pipeline_norm (StableDiffusionPipeline): The Stable Diffusion Pipeline for Norm Fine-Tuning
+#         sd_pipeline_bias (StableDiffusionPipeline): The Stable Diffusion Pipeline for Bias Fine-Tuning
+#         sd_pipeline_attention (StableDiffusionPipeline): The Stable Diffusion Pipeline for Attention Fine-Tuning
+#         sd_pipeline_NBA (StableDiffusionPipeline): The Stable Diffusion Pipeline for NBA Fine-Tuning
+#         sd_pipeline_difffit (StableDiffusionPipeline): The Stable Diffusion Pipeline for Difffit Fine-Tuning
+#     """
+
+#     # Dictionary containing the path to the best trained models for each PEFT type
+#     MODEL_PATH_DICT = {
+#         "full": "full_diffusion_pytorch_model.safetensors",
+#         "norm": "norm_diffusion_pytorch_model.safetensors",
+#         "bias": "bias_diffusion_pytorch_model.safetensors",
+#         "attention": "attention_diffusion_pytorch_model.safetensors",
+#         "norm_bias_attention": "norm_bias_attention_diffusion_pytorch_model.safetensors",
+#         "difffit": "difffit_diffusion_pytorch_model.safetensors",
+#     }
+
+#     device = "0"
+#     cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu"
+
+#     # Full FT
+#     unet_pretraining_type = "full"
+#     print("Loading Pipeline for Full Fine-Tuning")
+#     sd_pipeline_full = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     # Norm
+#     unet_pretraining_type = "norm"
+#     print("Loading Pipeline for Norm Fine-Tuning")
+#     sd_pipeline_norm = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     # bias
+#     unet_pretraining_type = "bias"
+#     print("Loading Pipeline for Bias Fine-Tuning")
+#     sd_pipeline_bias = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     # attention
+#     unet_pretraining_type = "attention"
+#     print("Loading Pipeline for Attention Fine-Tuning")
+#     sd_pipeline_attention = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     # NBA
+#     unet_pretraining_type = "norm_bias_attention"
+#     print("Loading Pipeline for NBA Fine-Tuning")
+#     sd_pipeline_NBA = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     # difffit
+#     unet_pretraining_type = "difffit"
+#     print("Loading Pipeline for Difffit Fine-Tuning")
+#     sd_pipeline_difffit = loadSDModel(
+#         unet_pretraining_type=unet_pretraining_type,
+#         exp_path=MODEL_PATH_DICT[unet_pretraining_type],
+#         cuda_device=cuda_device,
+#     )
+
+#     return (
+#         sd_pipeline_full,
+#         sd_pipeline_norm,
+#         sd_pipeline_bias,
+#         sd_pipeline_attention,
+#         sd_pipeline_NBA,
+#         sd_pipeline_difffit,
+#     )
 
 
 # LOAD ALL PIPELINES FIRST AND CACHE THEM
@@ -235,7 +235,6 @@ def predict(
     print("Loading Pipeline for {} Fine-Tuning".format(unet_pretraining_type))
     sd_pipeline_norm = loadSDModel(
         unet_pretraining_type=unet_pretraining_type,
-        exp_path=MODEL_PATH_DICT[unet_pretraining_type],
         cuda_device=cuda_device,
     )