more prompt

app.py

@@ -149,29 +149,43 @@ def update_sr_prompt(model_name):
         return "F-actin of COS-7"
     return "" # 或者返回一个默认值
 
+PROMPT_TO_MODEL_MAP = {}
+current_t2i_unet_path = None
 def load_all_prompts():
-    prompt_files = [
-        r"prompts/basic_prompts.json",
-        r"prompts/others_prompts.json",
-        r"prompts/hpa_prompts.json"
+    global PROMPT_TO_MODEL_MAP
+    categories = [
+        {
+            "file": "prompts/basic_prompts.json", 
+            "model": f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
+        },
+        {
+            "file": "prompts/others_prompts.json", 
+            "model": f"{MODELS_ROOT_DIR}/FluoGen-demo-test-ckpts/FULL-checkpoint-275000"
+        },
+        {
+            "file": "prompts/hpa_prompts.json", 
+            "model": f"{MODELS_ROOT_DIR}/FluoGen-demo-test-ckpts/HPA-checkpoint-40000"
+        }
     ]
     
-    combined_prompts = []
-    for file_path in prompt_files:
+    combined_prompts = [] 
+    for cat in categories:
+        file_path = cat["file"]
+        model_path = cat["model"]
         try:
             if os.path.exists(file_path):
                 with open(file_path, "r", encoding="utf-8") as f:
                     data = json.load(f)
                     if isinstance(data, list):
                         combined_prompts.extend(data)
-                        print(f"✓ Loaded {len(data)} prompts from: {os.path.basename(file_path)}")
-            else:
-                print(f"⚠ Warning: File not found: {file_path}")
+                        for p in data:
+                            PROMPT_TO_MODEL_MAP[p] = model_path
+                        print(f"✓ Loaded {len(data)} prompts from {file_path}")
         except Exception as e:
             print(f"✗ Error loading {file_path}: {e}")
 
     if not combined_prompts:
-        return ["F-actin of COS-7", "ER of COS-7", "Mitochondria of BPAE"]
+        return ["F-actin of COS-7", "ER of COS-7"]
     return combined_prompts
 T2I_PROMPTS = load_all_prompts()
 
@@ -186,6 +200,7 @@ try:
     t2i_tokenizer = CLIPTokenizer.from_pretrained(T2I_PRETRAINED_MODEL_PATH, subfolder="tokenizer")
     t2i_pipe = DDPMPipeline(unet=t2i_unet, scheduler=t2i_noise_scheduler, text_encoder=t2i_text_encoder, tokenizer=t2i_tokenizer)
     t2i_pipe.to(DEVICE)
+    current_t2i_unet_path = T2I_UNET_PATH
     print("✓ Text-to-Image model loaded successfully!")
 except Exception as e:
     print(f"!!!!!! FATAL: Text-to-Image Model Loading Failed !!!!!!\nError: {e}")
@@ -217,9 +232,40 @@ def swap_controlnet(pipe, target_path):
             raise gr.Error(f"Failed to load ControlNet model '{target_path}'. Error: {e}")
     return pipe
 
+def swap_t2i_unet(pipe, target_unet_path):
+    global current_t2i_unet_path
+    target_unet_path = os.path.normpath(target_unet_path)
+    if current_t2i_unet_path is None or os.path.normpath(current_t2i_unet_path) != target_unet_path:
+        print(f"🔄 Swapping T2I UNet to: {target_unet_path}")
+        try:
+            new_unet = UNet2DModel.from_pretrained(target_unet_path, subfolder="unet").to(DEVICE)
+            pipe.unet = new_unet
+            current_t2i_unet_path = target_unet_path
+            print("✅ UNet swapped successfully.")
+        except Exception as e:
+            raise gr.Error(f"Failed to load UNet from {target_unet_path}. Error: {e}")
+    return pipe
+
+# def generate_t2i(prompt, num_inference_steps):
+#     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
+#     print(f"\nTask started... | Prompt: '{prompt}' | Steps: {num_inference_steps}")
+#     image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
+#     generated_image = numpy_to_pil(image_np)
+#     print("✓ Image generated")
+#     if SAVE_EXAMPLES:
+#         example_filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, sanitize_prompt_for_filename(prompt))0
+#         if not os.path.exists(example_filepath):
+#             generated_image.save(example_filepath); print(f"✓ New T2I example saved: {example_filepath}")
+#     return generated_image
 def generate_t2i(prompt, num_inference_steps):
+    global t2i_pipe
     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
-    print(f"\nTask started... | Prompt: '{prompt}' | Steps: {num_inference_steps}")
+    target_model_path = PROMPT_TO_MODEL_MAP.get(prompt)
+    if target_model_path:
+        t2i_pipe = swap_t2i_unet(t2i_pipe, target_model_path)
+    else:
+        print(f"⚠️ Warning: No specific model mapped for '{prompt}', using current weights.")
+    print(f"\n🚀 Task started... | Prompt: '{prompt}' | Model: {current_t2i_unet_path}")
     image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
     generated_image = numpy_to_pil(image_np)
     print("✓ Image generated")
@@ -577,7 +623,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     # sr_prompt_input = gr.Textbox(label="Prompt (e.g., structure name)", value="CCPs of COS-7")
                     sr_prompt_input = gr.Textbox(
                         label="Prompt", 
-                        value="CCPs of COS-7", # 初始值根据你的默认选择设定
+                        value="F-actin of COS-7", # 初始值根据你的默认选择设定
                         interactive=False 
                     )
                     sr_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")