info labels removed

app.py

@@ -12,12 +12,14 @@ import yaml
 
 warnings.filterwarnings("ignore")
 
+################################################################################
+
+
+# Define the default parameters
 OUTPUT_DIR = "OUTPUT"
 cuda_device = 1
 device = f"cuda:{cuda_device}" if torch.cuda.is_available() else "cpu"
 
-print("DEVICE: ", device)
-
 TITLE = "Demo for Generating Chest X-rays using Diferent Parameter-Efficient Fine-Tuned Stable Diffusion Pipelines"
 INFO_ABOUT_TEXT_PROMPT = "INFO_ABOUT_TEXT_PROMPT"
 INFO_ABOUT_GUIDANCE_SCALE = "INFO_ABOUT_GUIDANCE_SCALE"
@@ -31,7 +33,7 @@ EXAMPLE_TEXT_PROMPTS = [
     "No radiographic evidence for acute cardiopulmonary process",
 ]
 
-
+################################################################################
 def load_adapted_unet(unet_pretraining_type, pipe):
 
     """
@@ -46,7 +48,7 @@ def load_adapted_unet(unet_pretraining_type, pipe):
     """
 
     sd_folder_path = "runwayml/stable-diffusion-v1-5"
-    exp_path = ''
+    exp_path = ""
 
     if unet_pretraining_type == "freeze":
         pass
@@ -71,7 +73,9 @@ def load_adapted_unet(unet_pretraining_type, pipe):
     else:
         # 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")
+        state_dict = load_file(
+            unet_pretraining_type + "_" + "diffusion_pytorch_model.safetensors"
+        )
         print(pipe.unet.load_state_dict(state_dict, strict=False))
 
 
@@ -97,18 +101,20 @@ def loadSDModel(unet_pretraining_type, cuda_device):
 
     return pipe
 
+
 def _predict_using_default_params():
 
-    
     # Defining the default parameters
-    unet_pretraining_type = 'full'
-    input_text = 'No acute cardiopulmonary abnormality.'
+    unet_pretraining_type = "full"
+    input_text = "No acute cardiopulmonary abnormality."
     guidance_scale = 4
     num_inference_steps = 75
-    device = '0'
-    OUTPUT_DIR = 'OUTPUT'
+    device = "0"
+    OUTPUT_DIR = "OUTPUT"
 
-    BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(unet_pretraining_type)
+    BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(
+        unet_pretraining_type
+    )
     NUM_TUNABLE_PARAMS = {
         "full": 86,
         "attention": 26.7,
@@ -150,9 +156,9 @@ def _predict_using_default_params():
 
     print(df)
 
-    df = df[df["Fine-Tuning Strategy"].isin(["full", unet_pretraining_type])].reset_index(
-        drop=True
-    )
+    df = df[
+        df["Fine-Tuning Strategy"].isin(["full", unet_pretraining_type])
+    ].reset_index(drop=True)
 
     bar_plot = gr.BarPlot(
         value=df,
@@ -177,8 +183,29 @@ def predict(
     OUTPUT_DIR="OUTPUT",
 ):
 
+    """
+    Generates a Chest X-ray using the selected PEFT Type, input text prompt, guidance scale, and number of inference steps
+
+    Parameters:
+        unet_pretraining_type (str): The type of PEFT to use for generating the X-ray
+        input_text (str): The text prompt to use for generating the X-ray
+        guidance_scale (int): The guidance scale to use for generating the X-ray
+        num_inference_steps (int): The number of inference steps to use for generating the X-ray
+        device (str): The CUDA device to use for generating the X-ray
+        OUTPUT_DIR (str): The output directory to save the generated X-ray
+
+    Returns:
+        result_pil_image (PIL.Image): The generated X-ray image
+        bar_plot (gr.BarPlot): The number of tunable parameters for the selected PEFT Type
+    """
+
+    # Run the _predict_using_default_params() function to generate a defualt X-ray output
+    # result_pil_image, bar_plot = _predict_using_default_params()
+
     try:
-        BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(unet_pretraining_type)
+        BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(
+            unet_pretraining_type
+        )
         NUM_TUNABLE_PARAMS = {
             "full": 86,
             "attention": 26.7,
@@ -191,7 +218,7 @@ def predict(
         }
 
         cuda_device = f"cuda:{device}" if torch.cuda.is_available() else "cpu"
-        
+
         print("Loading Pipeline for {} Fine-Tuning".format(unet_pretraining_type))
         sd_pipeline = loadSDModel(
             unet_pretraining_type=unet_pretraining_type,
@@ -209,7 +236,6 @@ def predict(
         )
 
         result_pil_image = result_image["images"][0]
-        
 
         # Create a Bar Plot displaying the number of tunable parameters for the selected PEFT Type
         df = pd.DataFrame(
@@ -221,9 +247,9 @@ def predict(
 
         print(df)
 
-        df = df[df["Fine-Tuning Strategy"].isin(["full", unet_pretraining_type])].reset_index(
-            drop=True
-        )
+        df = df[
+            df["Fine-Tuning Strategy"].isin(["full", unet_pretraining_type])
+        ].reset_index(drop=True)
 
         bar_plot = gr.BarPlot(
             value=df,
@@ -263,14 +289,15 @@ iface = gr.Interface(
             label="PEFT Type",
         ),
         gr.Dropdown(
-            EXAMPLE_TEXT_PROMPTS, info=INFO_ABOUT_TEXT_PROMPT, label="Input Text", value=EXAMPLE_TEXT_PROMPTS[0]
+            EXAMPLE_TEXT_PROMPTS,
+            label="Input Text",
+            value=EXAMPLE_TEXT_PROMPTS[0],
         ),
         gr.Slider(
             minimum=1,
             maximum=10,
             value=4,
             step=1,
-            info=INFO_ABOUT_GUIDANCE_SCALE,
             label="Guidance Scale",
         ),
         gr.Slider(
@@ -278,7 +305,6 @@ iface = gr.Interface(
             maximum=100,
             value=75,
             step=1,
-            info=INFO_ABOUT_INFERENCE_STEPS,
             label="Num Inference Steps",
         ),
     ],