Update app.py

app.py

@@ -245,22 +245,50 @@ if not torch.cuda.is_available():
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+# Define available models
+AVAILABLE_MODELS = {
+    "NaixlMmmmix v50": "naixlMmmmix_v50.safetensors",
+    "Miaomiao": "miaomiao.safetensors",
+    # Add more models here as needed
+}
+
+# Global variable to store the loaded model and pipeline
+loaded_model = None
+pipe = None
+vae = None
+
+# Initialize VAE
 if torch.cuda.is_available():
     vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-    #token = os.environ.get("HF_TOKEN")  # 从环境变量读取令牌
-    model_path = hf_hub_download(
-        repo_id="Menyu/ModelFile",  # 模型仓库名称（非完整URL）
-        filename="naixlMmmmix_v50.safetensors" #,
-        #use_auth_token=token
-    )
-    pipe = StableDiffusionXLPipeline.from_single_file(
-        model_path,
-        vae=vae,
-        use_safetensors=True,
-        torch_dtype=torch.float16,
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    pipe.to("cuda")
+
+def load_model(model_name):
+    global pipe, loaded_model
+    
+    # If the requested model is already loaded, return the existing pipeline
+    if loaded_model == model_name and pipe is not None:
+        return pipe
+    
+    if torch.cuda.is_available():
+        model_filename = AVAILABLE_MODELS[model_name]
+        model_path = hf_hub_download(
+            repo_id="Menyu/ModelFile",
+            filename=model_filename
+        )
+        
+        pipe = StableDiffusionXLPipeline.from_single_file(
+            model_path,
+            vae=vae,
+            use_safetensors=True,
+            torch_dtype=torch.float16,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipe.to("cuda")
+        
+        # Update the loaded model name
+        loaded_model = model_name
+        
+        return pipe
+    return None
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -279,11 +307,18 @@ def infer(
     num_inference_steps: int = 30,
     randomize_seed: bool = True,
     use_resolution_binning: bool = True,
+    model_name: str = "NaixlMmmmix v50",
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Load the selected model
+    pipe = load_model(model_name)
+    if pipe is None:
+        return None, seed
+        
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator().manual_seed(seed)
-    # 初始化 Compel 实例
+    
+    # Initialize Compel instance
     compel = Compel(
         tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
         text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
@@ -291,18 +326,17 @@ def infer(
         requires_pooled=[False, True],
         truncate_long_prompts=False
     )
+    
     # Store the original prompt before processing
     original_prompt_text = prompt
     
-    # 在 infer 函数中调用 get_embed_new
     if not use_negative_prompt:
         negative_prompt = ""
     
     processed_prompt = get_embed_new(prompt, pipe, compel, only_convert_string=True)
     processed_negative_prompt = get_embed_new(negative_prompt, pipe, compel, only_convert_string=True)
-    conditioning, pooled = compel([processed_prompt, processed_negative_prompt]) # 必须同时处理来保证长度相等
+    conditioning, pooled = compel([processed_prompt, processed_negative_prompt]) # Process both to ensure equal length
     
-    # 在调用 pipe 时，使用新的参数名称（确保参数名称正确）
     image = pipe(
         prompt_embeds=conditioning[0:1],
         pooled_prompt_embeds=pooled[0:1],
@@ -326,7 +360,7 @@ def infer(
         "height": height,
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
-        "model": "naixlMmmmix_v50",
+        "model": model_name,
         "use_resolution_binning": use_resolution_binning,
         "PreUrl": "https://huggingface.co/spaces/Menyu/NaixlMix"
     }
@@ -348,7 +382,7 @@ h1{text-align:center}
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown("""# 梦羽的模型生成器
-        ### 快速生成NaixlMmmmix v50模型的图片""")
+        ### 快速生成AI模型的图片""")
     with gr.Group():
         with gr.Row():
             prompt = gr.Text(
@@ -361,6 +395,11 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("生成", scale=0, variant="primary")
         result = gr.Image(label="Result", show_label=False, format="png")
     with gr.Accordion("高级选项", open=False):
+        model_selector = gr.Dropdown(
+            label="选择模型",
+            choices=list(AVAILABLE_MODELS.keys()),
+            value="NaixlMmmmix v50"
+        )
         with gr.Row():
             use_negative_prompt = gr.Checkbox(label="使用反向词条", value=True)
             negative_prompt = gr.Text(
@@ -414,7 +453,7 @@ with gr.Blocks(css=css) as demo:
         examples=examples,
         inputs=prompt,
         outputs=[result, seed],
-        fn=infer
+        fn=lambda x: infer(x, model_name=model_selector.value)
     )
 
     use_negative_prompt.change(
@@ -436,6 +475,8 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed,
+            gr.Checkbox(value=True, visible=False, label="use_resolution_binning"),  # Hidden default value
+            model_selector,
         ],
         outputs=[result, seed],
     )