Update app.py

app.py

@@ -10,7 +10,14 @@ from insightface.utils import face_align
 import gradio as gr
 import cv2
 
-base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
+base_model_paths = {
+    "Realistic_Vision_V4.0_noVAE": "SG161222/Realistic_Vision_V4.0_noVAE",
+    "Realistic_Vision_V6.0_B1_noVAE": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
+    "Deliberate": "Yntec/Deliberate",
+    "Deliberate2": "Yntec/Deliberate2",
+    "pony-diffusion": "AstraliteHeart/pony-diffusion"
+}
+
 vae_model_path = "stabilityai/sd-vae-ft-mse"
 image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
 ip_ckpt = hf_hub_download(repo_id="h94/IP-Adapter-FaceID", filename="ip-adapter-faceid_sd15.bin", repo_type="model")
@@ -32,20 +39,20 @@ noise_scheduler = DDIMScheduler(
     steps_offset=1,
 )
 vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch.float16)
-pipe = StableDiffusionPipeline.from_pretrained(
-    base_model_path,
-    torch_dtype=torch.float16,
-    scheduler=noise_scheduler,
-    vae=vae,
-    feature_extractor=safety_feature_extractor,
-    safety_checker=None  # <--- Disable safety checker
-).to(device)
 
-#pipe.load_lora_weights("h94/IP-Adapter-FaceID", weight_name="ip-adapter-faceid-plusv2_sd15_lora.safetensors")
-#pipe.fuse_lora()
+def load_model(base_model_path):
+    pipe = StableDiffusionPipeline.from_pretrained(
+        base_model_path,
+        torch_dtype=torch.float16,
+        scheduler=noise_scheduler,
+        vae=vae,
+        feature_extractor=safety_feature_extractor,
+        safety_checker=None  # <--- Disable safety checker
+    ).to(device)
+    return pipe
 
-ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
-ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
+ip_model = None
+ip_model_plus = None
 
 app = FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640))
@@ -53,7 +60,13 @@ app.prepare(ctx_id=0, det_size=(640, 640))
 cv2.setNumThreads(1)
 
 @spaces.GPU(enable_queue=True)
-def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, progress=gr.Progress(track_tqdm=True)):
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, base_model, progress=gr.Progress(track_tqdm=True)):
+    global ip_model, ip_model_plus
+    base_model_path = base_model_paths[base_model]
+    pipe = load_model(base_model_path)
+    ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
+    ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
+
     faceid_all_embeds = []
     first_iteration = True
     for image in images:
@@ -95,10 +108,12 @@ def swap_to_gallery(images):
 
 def remove_back_to_files():
     return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+
 css = '''
 h1{margin-bottom: 0 !important}
 footer{display:none !important}
 '''
+
 with gr.Blocks(css=css) as demo:
     gr.Markdown("")
     gr.Markdown("")
@@ -116,6 +131,7 @@ with gr.Blocks(css=css) as demo:
                        placeholder="A photo of a [man/woman/person]...")
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality")
             style = gr.Radio(label="Generation type", info="For stylized try prompts like 'a watercolor painting of a woman'", choices=["Photorealistic", "Stylized"], value="Photorealistic")
+            base_model = gr.Dropdown(label="Base Model", choices=list(base_model_paths.keys()), value="Realistic_Vision_V4.0_noVAE")
             submit = gr.Button("Submit")
             with gr.Accordion(open=False, label="Advanced Options"):
                 preserve = gr.Checkbox(label="Preserve Face Structure", info="Higher quality, less versatility (the face structure of your first photo will be preserved). Unchecking this will use the v1 model.", value=True)
@@ -130,7 +146,7 @@ with gr.Blocks(css=css) as demo:
         files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
         remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
         submit.click(fn=generate_image,
-                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts],
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts, base_model],
                     outputs=gallery)
     
     gr.Markdown("")