_predict_using_default_params function added for failing cases

app.py

@@ -97,15 +97,16 @@ def loadSDModel(unet_pretraining_type, cuda_device):
 
     return pipe
 
+def _predict_using_default_params():
 
-def predict(
-    unet_pretraining_type,
-    input_text,
-    guidance_scale=4,
-    num_inference_steps=75,
-    device="0",
-    OUTPUT_DIR="OUTPUT",
-):
+    
+    # Defining the default parameters
+    unet_pretraining_type = 'full'
+    input_text = 'No acute cardiopulmonary abnormality.'
+    guidance_scale = 4
+    num_inference_steps = 75
+    device = '0'
+    OUTPUT_DIR = 'OUTPUT'
 
     BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(unet_pretraining_type)
     NUM_TUNABLE_PARAMS = {
@@ -120,7 +121,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,
@@ -138,7 +139,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(
@@ -168,6 +168,80 @@ def predict(
     return result_pil_image, bar_plot
 
 
+def predict(
+    unet_pretraining_type,
+    input_text,
+    guidance_scale=4,
+    num_inference_steps=75,
+    device="0",
+    OUTPUT_DIR="OUTPUT",
+):
+
+    try:
+        BARPLOT_TITLE = "Tunable Parameters for {} Fine-Tuning".format(unet_pretraining_type)
+        NUM_TUNABLE_PARAMS = {
+            "full": 86,
+            "attention": 26.7,
+            "bias": 0.343,
+            "norm": 0.2,
+            "norm_bias_attention": 26.7,
+            "lorav2": 0.8,
+            "svdiff": 0.222,
+            "difffit": 0.581,
+        }
+
+        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,
+            cuda_device=cuda_device,
+        )
+
+        sd_pipeline.to(cuda_device)
+
+        result_image = sd_pipeline(
+            prompt=input_text,
+            height=224,
+            width=224,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+        )
+
+        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(
+            {
+                "Fine-Tuning Strategy": list(NUM_TUNABLE_PARAMS.keys()),
+                "Number of Tunable Parameters": list(NUM_TUNABLE_PARAMS.values()),
+            }
+        )
+
+        print(df)
+
+        df = df[df["Fine-Tuning Strategy"].isin(["full", unet_pretraining_type])].reset_index(
+            drop=True
+        )
+
+        bar_plot = gr.BarPlot(
+            value=df,
+            x="Fine-Tuning Strategy",
+            y="Number of Tunable Parameters",
+            title=BARPLOT_TITLE,
+            vertical=True,
+            height=300,
+            width=300,
+            interactive=True,
+        )
+
+        return result_pil_image, bar_plot
+
+    except:
+        return _predict_using_default_params()
+
+
 # Create a Gradio interface
 """
 Input Parameters: