Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 import spaces
-from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, StableDiffusionXLPipeline
+from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
 from transformers import AutoFeatureExtractor
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 from ip_adapter.ip_adapter_faceid import IPAdapterFaceID, IPAdapterFaceIDPlus
@@ -9,7 +9,6 @@ from insightface.app import FaceAnalysis
 from insightface.utils import face_align
 import gradio as gr
 import cv2
-import os
 
 base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE"
 vae_model_path = "stabilityai/sd-vae-ft-mse"
@@ -42,6 +41,9 @@ pipe = StableDiffusionPipeline.from_pretrained(
     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()
+
 ip_model = IPAdapterFaceID(pipe, ip_ckpt, device)
 ip_model_plus = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_plus_ckpt, device)
 
@@ -50,22 +52,8 @@ app.prepare(ctx_id=0, det_size=(640, 640))
 
 cv2.setNumThreads(1)
 
-# Download the SDXL model files
-ckpt_dir_pony = snapshot_download(repo_id="John6666/pony-realism-v21main-sdxl")
-ckpt_dir_cyber = snapshot_download(repo_id="John6666/cyberrealistic-pony-v61-sdxl")
-ckpt_dir_stallion = snapshot_download(repo_id="John6666/stallion-dreams-pony-realistic-v1-sdxl")
-
-# Load the SDXL models
-pipe_pony = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_pony, torch_dtype=torch.float16)
-pipe_cyber = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_cyber, torch_dtype=torch.float16)
-pipe_stallion = StableDiffusionXLPipeline.from_pretrained(ckpt_dir_stallion, torch_dtype=torch.float16)
-
-pipe_pony.to(device)
-pipe_cyber.to(device)
-pipe_stallion.to(device)
-
 @spaces.GPU(enable_queue=True)
-def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, model_choice, progress=gr.Progress(track_tqdm=True)):
+def generate_image(images, prompt, negative_prompt, preserve_face_structure, face_strength, likeness_strength, nfaa_negative_prompt, progress=gr.Progress(track_tqdm=True)):
     faceid_all_embeds = []
     first_iteration = True
     for image in images:
@@ -81,24 +69,17 @@ def generate_image(images, prompt, negative_prompt, preserve_face_structure, fac
     
     total_negative_prompt = f"{negative_prompt} {nfaa_negative_prompt}"
     
-    if model_choice == "Pony Realism v21":
-        pipe = pipe_pony
-    elif model_choice == "Cyber Realistic Pony v61":
-        pipe = pipe_cyber
-    else:  # "Stallion Dreams Pony Realistic v1"
-        pipe = pipe_stallion
-    
     if(not preserve_face_structure):
         print("Generating normal")
         image = ip_model.generate(
             prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
-            scale=likeness_strength, width=512, height=512, num_inference_steps=30, pipe=pipe
+            scale=likeness_strength, width=512, height=512, num_inference_steps=30
         )
     else:
         print("Generating plus")
         image = ip_model_plus.generate(
             prompt=prompt, negative_prompt=total_negative_prompt, faceid_embeds=average_embedding,
-            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=512, height=512, num_inference_steps=30, pipe=pipe
+            scale=likeness_strength, face_image=face_image, shortcut=True, s_scale=face_strength, width=512, height=512, num_inference_steps=30
         )
     print(image)
     return image
@@ -114,12 +95,10 @@ 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("")
@@ -137,11 +116,6 @@ 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")
-            model_choice = gr.Dropdown(
-                ["Pony Realism v21", "Cyber Realistic Pony v61", "Stallion Dreams Pony Realistic v1"],
-                label="Model Choice",
-                value="Pony Realism v21"
-            )
             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)
@@ -156,7 +130,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, model_choice],
+                    inputs=[files,prompt,negative_prompt,preserve, face_strength, likeness_strength, nfaa_negative_prompts],
                     outputs=gallery)
     
     gr.Markdown("")