Update app.py

app.py

@@ -11,7 +11,7 @@ from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, Auto
 from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
 from controlnet_aux import PidiNetDetector, HEDdetector
 from diffusers.utils import load_image
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, snapshot_download
 from pathlib import Path
 from PIL import Image, ImageOps
 import cv2
@@ -124,41 +124,37 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
 
+# Download the 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")
 
-controlnet = ControlNetModel.from_pretrained(
-    "xinsir/controlnet-union-sdxl-1.0",
-    torch_dtype=torch.float16
-)
-controlnet_canny = ControlNetModel.from_pretrained(
-    "xinsir/controlnet-union-sdxl-1.0",
-    torch_dtype=torch.float16
+# Load the models
+vae_pony = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_pony, "vae"), torch_dtype=torch.float16)
+vae_cyber = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_cyber, "vae"), torch_dtype=torch.float16)
+vae_stallion = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_stallion, "vae"), torch_dtype=torch.float16)
+
+controlnet_pony = ControlNetModel.from_pretrained("xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16)
+controlnet_cyber = ControlNetModel.from_pretrained("xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16)
+controlnet_stallion = ControlNetModel.from_pretrained("xinsir/controlnet-union-sdxl-1.0", torch_dtype=torch.float16)
+
+pipe_pony = StableDiffusionXLControlNetPipeline.from_pretrained(
+    ckpt_dir_pony, controlnet=controlnet_pony, vae=vae_pony, torch_dtype=torch.float16, scheduler=eulera_scheduler
 )
-# when test with other base model, you need to change the vae also.
-vae = AutoencoderKL.from_pretrained("John6666/pony-realism-v21main-sdxl", subfolder="vae", torch_dtype=torch.float16)
-
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-    "John6666/pony-realism-v21main-sdxl",
-    controlnet=controlnet,
-    vae=vae,
-    torch_dtype=torch.float16,
-    scheduler=eulera_scheduler,
+pipe_cyber = StableDiffusionXLControlNetPipeline.from_pretrained(
+    ckpt_dir_cyber, controlnet=controlnet_cyber, vae=vae_cyber, torch_dtype=torch.float16, scheduler=eulera_scheduler
 )
-pipe.to(device)
-# Load model.
-pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
-    "John6666/pony-realism-v21main-sdxl",
-    controlnet=controlnet_canny,
-    vae=vae,
-    safety_checker=None,
-    torch_dtype=torch.float16,
-    scheduler=eulera_scheduler,
+pipe_stallion = StableDiffusionXLControlNetPipeline.from_pretrained(
+    ckpt_dir_stallion, controlnet=controlnet_stallion, vae=vae_stallion, torch_dtype=torch.float16, scheduler=eulera_scheduler
 )
-pipe_canny.to(device)
+
+pipe_pony.to(device)
+pipe_cyber.to(device)
+pipe_stallion.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
@@ -189,6 +185,7 @@ def run(
     image: dict,
     prompt: str,
     negative_prompt: str,
+    model_choice: str,  # Add this new input
     style_name: str = DEFAULT_STYLE_NAME,
     num_steps: int = 25,
     guidance_scale: float = 5,
@@ -233,8 +230,16 @@ def run(
 
     generator = torch.Generator(device=device).manual_seed(seed)
     
+    # Select the appropriate pipe based on the model choice
+    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 use_canny:
-        out = pipe_canny(
+        out = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             image=image,
@@ -273,7 +278,11 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
             with gr.Group():
                 image = gr.ImageEditor(type="pil", label="Sketch your image or upload one", width=512, height=512)
                 prompt = gr.Textbox(label="Prompt")
-                style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                model_choice = gr.Dropdown(
+                    ["Pony Realism v21", "Cyber Realistic Pony v61", "Stallion Dreams Pony Realistic v1"],
+                    label="Model Choice",
+                    value="Pony Realism v21"
+                )
                 use_hed = gr.Checkbox(label="use HED detector", value=False, info="check this box if you upload an image and want to turn it to a sketch")
                 use_canny = gr.Checkbox(label="use Canny", value=False, info="check this to use ControlNet canny instead of scribble")
                 run_button = gr.Button("Run")
@@ -321,6 +330,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
         image,
         prompt,
         negative_prompt,
+        model_choice,  # Add this new input
         style,
         num_steps,
         guidance_scale,