Update app.py

app.py

@@ -268,7 +268,7 @@ if torch.cuda.is_available():
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
         print("Model RealVisXL_V5.0_Lightning Compiled!")
     
-    # Load new model (RealVisXL_V4.0)
+    # Load second model (RealVisXL_V4.0)
     pipe2 = StableDiffusionXLPipeline.from_pretrained(
         "SG161222/RealVisXL_V4.0",
         torch_dtype=dtype,
@@ -286,8 +286,27 @@ if torch.cuda.is_available():
     if USE_TORCH_COMPILE:
         pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
         print("Model RealVisXL_V4.0 Compiled!")
+    
+    # Load third model (Animagine XL 4.0)
+    pipe3 = StableDiffusionXLPipeline.from_pretrained(
+        "cagliostrolab/animagine-xl-4.0",
+        torch_dtype=dtype,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(device)
+    pipe3.text_encoder = pipe3.text_encoder.half()
+    
+    if ENABLE_CPU_OFFLOAD:
+        pipe3.enable_model_cpu_offload()
+    else:
+        pipe3.to(device)
+        print("Loaded Animagine XL 4.0 on Device!")
+    
+    if USE_TORCH_COMPILE:
+        pipe3.unet = torch.compile(pipe3.unet, mode="reduce-overhead", fullgraph=True)
+        print("Model Animagine XL 4.0 Compiled!")
 else:
-    # On CPU load both models in float32
+    # On CPU, load all models in float32
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "SG161222/RealVisXL_V5.0_Lightning",
         torch_dtype=dtype,
@@ -300,14 +319,21 @@ else:
         use_safetensors=True,
         add_watermarker=False,
     ).to(device)
+    pipe3 = StableDiffusionXLPipeline.from_pretrained(
+        "cagliostrolab/animagine-xl-4.0",
+        torch_dtype=dtype,
+        use_safetensors=True,
+        add_watermarker=False,
+    ).to(device)
     print("Running on CPU; models loaded in float32.")
 
 # A dictionary to easily choose the model based on selection.
 DEFAULT_MODEL = "RealVisXL_V5.0_Lightning"
-MODEL_CHOICES = [DEFAULT_MODEL, "RealVisXL_V4.0"]
+MODEL_CHOICES = [DEFAULT_MODEL, "RealVisXL_V4.0", "Animagine XL 4.0"]
 models = {
     "RealVisXL_V5.0_Lightning": pipe,
-    "RealVisXL_V4.0": pipe2
+    "RealVisXL_V4.0": pipe2,
+    "Animagine XL 4.0": pipe3
 }
 
 def save_image(img, path):