Update app.py

app.py

@@ -31,6 +31,18 @@ footer {
 }
 '''
 
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+controlnet = ControlNetModel.from_pretrained("MakiPan/controlnet-encoded-hands-130k", torch_dtype=torch.float16)
+
+repo_default = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
+repo_default.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
+repo_default.set_adapters(["base"], adapter_weights=[0.7])
+
+repo_customs = {
+    "Default": repo_default,
+    "Anime": StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False),
+}
+
 # Functions
 def save_image(img, seed):
     name = f"{seed}-{uuid.uuid4()}.png"
@@ -47,7 +59,7 @@ def get_seed(seed):
 @spaces.GPU(duration=30)
 def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATIVE_INPUT, model=DEFAULT_MODEL, height=DEFAULT_HEIGHT, width=DEFAULT_WIDTH, steps=1, guidance=0, number=1, seed=None):
 
-    repo = None
+    repo = repo_customs[model or "Default"]
     filter_input = filter_input or ""
     negative_input = negative_input or DEFAULT_NEGATIVE_INPUT
     seed = get_seed(seed)
@@ -55,17 +67,9 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     print(input, filter_input, negative_input, model, height, width, steps, guidance, number, seed)
 
     if model == "Anime":   
-        vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-        controlnet = ControlNetModel.from_pretrained("MakiPan/controlnet-encoded-hands-130k", torch_dtype=torch.float16)
-        repo = StableDiffusionXLPipeline.from_pretrained("cagliostrolab/animagine-xl-3.1", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
         steps = (not steps or steps < 0 and 24) or steps
         guidance = (not guidance or guidance < 0 and 7) or guidance
     else:
-        vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-        controlnet = ControlNetModel.from_pretrained("MakiPan/controlnet-encoded-hands-130k", torch_dtype=torch.float16)
-        repo = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
-        repo.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
-        repo.set_adapters(["base"], adapter_weights=[0.7])
         steps = (not steps or steps < 0 and 16) or steps
         guidance = (not guidance or guidance < 0 and 3) or guidance