Update app.py

app.py

@@ -40,13 +40,25 @@ def get_prediction(img):
 
 models = ['DDPM', 'Glide', 'Latent Diffusion', 'Palette', 'Stable Diffusion', 'VQ Diffusion', 'real', 'unseen_fake']
 
-local_model_path = snapshot_download(repo_id="ElBeh/ma_basemodel")
-model = keras.models.load_model(local_model_path)
+st.title("Fake Detection")
+st.divider()
 
+st.subheader("Modelvariant")
+variant = st.selectbox(
+    "Choose the model",
+    ("ResNet50v2-Basemodel", "ResNet50v2-Finetuned"),
+)
 
-st.title("Fake Detection")
+st.write("You selected model: ", variant)
 st.divider()
 
+if variant == "ResNet50v2-Basemodel":
+    local_model_path = snapshot_download(repo_id="ElBeh/ma_basemodel")
+else:
+    local_model_path = snapshot_download(repo_id="ElBeh/ma_finetuned_model")
+    
+model = keras.models.load_model(local_model_path)
+
 st.subheader("Image Preprocessing")
 crop = st.toggle("random crop")
 compress = st.toggle("jpeg compression")