Update app.py

app.py

@@ -133,15 +133,20 @@ eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/
 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")
+ckpt_dir_stable = snapshot_download(repo_id="stabilityai/stable-diffusion-xl-base-1.0")
+
 
 # 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)
+vae_stable = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_stable, "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)
+controlnet_stable = 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
@@ -152,10 +157,15 @@ pipe_cyber = StableDiffusionXLControlNetPipeline.from_pretrained(
 pipe_stallion = StableDiffusionXLControlNetPipeline.from_pretrained(
     ckpt_dir_stallion, controlnet=controlnet_stallion, vae=vae_stallion, torch_dtype=torch.float16, scheduler=eulera_scheduler
 )
+pipe_stable = StableDiffusionXLControlNetPipeline.from_pretrained(
+    ckpt_dir_stable, controlnet=controlnet_stable, vae=vae_stable, torch_dtype=torch.float16, scheduler=eulera_scheduler
+)
+
 
 pipe_pony.to(device)
 pipe_cyber.to(device)
 pipe_stallion.to(device)
+pipe_stable.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
@@ -232,12 +242,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":
+    if model_choice == "Stable Diffusion XL":
+        pipe = pipe_stable
+    elif model_choice == "Pony Realism v21":
         pipe = pipe_pony
     elif model_choice == "Cyber Realistic Pony v61":
         pipe = pipe_cyber
-    else:  # "Stallion Dreams Pony Realistic v1"
+    elif model_choice == "Stallion Dreams Pony Realistic v1":
         pipe = pipe_stallion
+    else:
+        raise ValueError("Invalid model choice")
     
     if use_canny:
         out = pipe(
@@ -281,7 +295,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
                 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"],
+                    ["Pony Realism v21", "Cyber Realistic Pony v61", "Stallion Dreams Pony Realistic v1", "Stable Diffusion XL"],
                     label="Model Choice",
                     value="Pony Realism v21"
                 )