Pseudo-color-6

app.py

@@ -49,6 +49,13 @@ WEIGHT_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 LOGO_PATH = "utils/logo2_transparent.png"
 SAVE_EXAMPLES = False
 
+# --- CSS for Times New Roman ---
+CUSTOM_CSS = """
+.gradio-container, .gradio-container * {
+    font-family: 'Times New Roman', Times, serif !important;
+}
+"""
+
 # --- Base directory for all models ---
 REPO_ID = "FluoGen-Group/FluoGen-demo-test-ckpts"
 MODELS_ROOT_DIR = snapshot_download(repo_id=REPO_ID, token=hf_token) 
@@ -303,24 +310,52 @@ def get_gallery_selection(evt: gr.SelectData):
 
 # --- Generation Functions ---
 @spaces.GPU(duration=120)
-def generate_t2i(prompt, num_inference_steps, current_color):
+def generate_t2i(prompt, num_inference_steps, num_images, current_color):
+    """
+    Generates multiple images for Text-to-Image and returns a gallery.
+    """
     global t2i_pipe
     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
     target_model_path = PROMPT_TO_MODEL_MAP.get(prompt, f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000")
     t2i_pipe = swap_t2i_unet(t2i_pipe, target_model_path)
     
-    print(f"\n🚀 T2I Task started... | Prompt: '{prompt}'")
-    image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
+    print(f"\n🚀 T2I Task started... | Prompt: '{prompt}' | Count: {num_images}")
     
-    raw_file_path = save_temp_tiff(image_np, prefix="t2i_raw")
-    display_image = apply_pseudocolor(image_np, current_color)
-    colorbar_img = generate_colorbar_preview(current_color)
+    generated_raw_list = []
+    generated_display_images = []
+    generated_raw_files = []
+    temp_dir = tempfile.mkdtemp()
     
-    if SAVE_EXAMPLES:
-        example_filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, sanitize_prompt_for_filename(prompt))
-        if not os.path.exists(example_filepath): display_image.save(example_filepath)
+    # Generate Batch
+    for i in range(int(num_images)):
+        # Generate single image
+        image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
+        generated_raw_list.append(image_np)
+        
+        # Save raw to temp
+        raw_name = f"t2i_sample_{i+1}.tif"
+        raw_path = os.path.join(temp_dir, raw_name)
+        save_data = image_np.astype(np.float32) if image_np.dtype == np.float16 else image_np
+        tifffile.imwrite(raw_path, save_data)
+        generated_raw_files.append(raw_path)
+        
+        # Create display version
+        generated_display_images.append(apply_pseudocolor(image_np, current_color))
+
+        # Save first image to examples if needed
+        if SAVE_EXAMPLES and i == 0:
+            example_filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, sanitize_prompt_for_filename(prompt))
+            if not os.path.exists(example_filepath): generated_display_images[0].save(example_filepath)
+
+    # Zip raw files
+    zip_filename = os.path.join(temp_dir, "raw_output_images.zip")
+    with zipfile.ZipFile(zip_filename, 'w') as zipf:
+        for file in generated_raw_files: zipf.write(file, os.path.basename(file))
+
+    colorbar_img = generate_colorbar_preview(current_color)
     
-    return display_image, raw_file_path, image_np, colorbar_img
+    # Return: Gallery List, Zip Path, Raw State List, Colorbar
+    return generated_display_images, zip_filename, generated_raw_list, colorbar_img
 
 @spaces.GPU(duration=120)
 def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed, current_color):
@@ -491,7 +526,7 @@ seg_examples = load_examples(SEG_EXAMPLE_IMG_DIR)
 cls_examples = load_examples(CLS_EXAMPLE_IMG_DIR)
 
 # --- UI Builders ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
     with gr.Row():
         gr.Image(value=LOGO_PATH, width=300, height=200, container=False, interactive=False, show_download_button=False, show_fullscreen_button=False)
         gr.Markdown(f"# {MODEL_TITLE}\n{MODEL_DESCRIPTION}")
@@ -499,25 +534,29 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Tabs():
         # --- TAB 1: Text-to-Image ---
         with gr.Tab("Text-to-Image Generation", id="txt2img"):
-            t2i_raw_state = gr.State(None)
+            t2i_raw_state = gr.State(None) # Stores list of arrays
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
                     t2i_prompt = gr.Dropdown(choices=T2I_PROMPTS, value=T2I_PROMPTS[0], label="Search or Type a Prompt", filterable=True, allow_custom_value=True)
                     t2i_steps = gr.Slider(10, 200, 50, step=1, label="Inference Steps")
+                    # Added: Number of Images Slider
+                    t2i_num_images = gr.Slider(1, 9, 4, step=1, label="Number of Images")
                     t2i_btn = gr.Button("Generate", variant="primary")
                 with gr.Column(scale=2):
-                    t2i_out = gr.Image(label="Generated Image", type="pil", interactive=False, show_download_button=False)
+                    # Changed: Image to Gallery
+                    t2i_gallery_out = gr.Gallery(label="Generated Images", columns=3, height="auto")
                     with gr.Row(equal_height=True):
                         with gr.Column(scale=2):
                             t2i_color = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (Adjust after generation)")
                         with gr.Column(scale=2):
                             t2i_colorbar = gr.Image(label="Colorbar", show_label=False, container=False, height=40, show_download_button=False, interactive=False)
                         with gr.Column(scale=1):
-                            t2i_dl = gr.DownloadButton(label="Download Raw (.tif)")
+                            t2i_dl = gr.DownloadButton(label="Download All (.zip)")
             t2i_gal = gr.Gallery(value=t2i_examples, label="Examples", columns=6, height="auto")
 
-            t2i_btn.click(generate_t2i, [t2i_prompt, t2i_steps, t2i_color], [t2i_out, t2i_dl, t2i_raw_state, t2i_colorbar])
-            t2i_color.change(update_single_image_color, [t2i_raw_state, t2i_color], [t2i_out, t2i_colorbar])
+            t2i_btn.click(generate_t2i, [t2i_prompt, t2i_steps, t2i_num_images, t2i_color], [t2i_gallery_out, t2i_dl, t2i_raw_state, t2i_colorbar])
+            # Reuse update_gallery_color since state is now a list
+            t2i_color.change(update_gallery_color, [t2i_raw_state, t2i_color], [t2i_gallery_out, t2i_colorbar])
             t2i_gal.select(fn=get_gallery_selection, inputs=None, outputs=t2i_prompt)
 
         # --- TAB 2: Super-Resolution ---
@@ -594,7 +633,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
                     m2i_file = gr.File(label="Upload Mask (.tif)", file_types=['.tif', '.tiff'])
-                    m2i_type = gr.Textbox(label="Cell Type", placeholder="e.g., HeLa")
+                    # Changed: Default value to HeLa
+                    m2i_type = gr.Textbox(label="Cell Type", value="HeLa", placeholder="e.g., HeLa")
                     m2i_num = gr.Slider(1, 10, 5, step=1, label="Count")
                     m2i_steps = gr.Slider(5, 50, 10, step=1, label="Steps")
                     m2i_seed = gr.Number(label="Seed", value=42)