feat: adding support for IN+

app.py

@@ -14,6 +14,7 @@ from gradio_log import Log
 # --- InstaNovo Imports ---
 try:
     from instanovo.transformer.model import InstaNovo
+    from instanovo.diffusion.multinomial_diffusion import InstaNovoPlus
     from instanovo.utils import SpectrumDataFrame, ResidueSet, Metrics
     from instanovo.transformer.dataset import SpectrumDataset, collate_batch
     from instanovo.inference import (
@@ -23,29 +24,38 @@ try:
         ScoredSequence,
         Decoder,
     )
-    from instanovo.constants import MASS_SCALE, MAX_MASS
+    from instanovo.inference.diffusion import DiffusionDecoder
+    from instanovo.constants import (
+        MASS_SCALE,
+        MAX_MASS,
+        DIFFUSION_START_STEP,
+    )
     from torch.utils.data import DataLoader
+    import torch.nn.functional as F # For padding
 except ImportError as e:
     raise ImportError(f"Failed to import InstaNovo components: {e}")
 
 # --- Configuration ---
-MODEL_ID = "instanovo-v1.1.0"  # Use the desired pretrained model ID
+TRANSFORMER_MODEL_ID = "instanovo-v1.1.0"
+DIFFUSION_MODEL_ID = "instanovoplus-v1.1.0-alpha"
 KNAPSACK_DIR = Path("./knapsack_cache")
 DEFAULT_CONFIG_PATH = Path(
     "./configs/inference/default.yaml"
-)  # Assuming instanovo installs configs locally relative to execution
+)
 
 # Determine device
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-FP16 = DEVICE == "cuda"  # Enable FP16 only on CUDA
+FP16 = DEVICE == "cuda"
 
-# --- Global Variables (Load Model and Knapsack Once) ---
+# --- Global Variables (Load Models and Knapsack Once) ---
 MODEL: InstaNovo | None = None
-KNAPSACK: Knapsack | None = None
 MODEL_CONFIG: DictConfig | None = None
+MODEL_PLUS: InstaNovoPlus | None = None
+MODEL_PLUS_CONFIG: DictConfig | None = None
+KNAPSACK: Knapsack | None = None
 RESIDUE_SET: ResidueSet | None = None
 
-# --- Assets --- 
+# --- Assets ---
 gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
 
 # Create gradio temporary directory
@@ -57,141 +67,165 @@ if not temp_dir.exists():
 log_file = "/tmp/instanovo_gradio_log.txt"
 Path(log_file).touch()
 
-logger = logging.getLogger("instanovo")
+logger = logging.getLogger("instanovo_gradio")
 logger.setLevel(logging.INFO)
-file_handler = logging.FileHandler(log_file)
-file_handler.setLevel(logging.INFO)
-formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
-file_handler.setFormatter(formatter)
-logger.addHandler(file_handler)
+if not logger.handlers:
+    file_handler = logging.FileHandler(log_file)
+    file_handler.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+    file_handler.setFormatter(formatter)
+    logger.addHandler(file_handler)
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+
+def load_models_and_knapsack():
+    """Loads the InstaNovo models and generates/loads the knapsack."""
+    global MODEL, KNAPSACK, MODEL_CONFIG, RESIDUE_SET, MODEL_PLUS, MODEL_PLUS_CONFIG
+    models_loaded = MODEL is not None and MODEL_PLUS is not None
+    if models_loaded:
+        logger.info("Models already loaded.")
+        # Still check knapsack if not loaded
+        if KNAPSACK is None:
+             logger.info("Models loaded, but knapsack needs loading/generation.")
+        else:
+             return # All loaded
 
+    # --- Load Transformer Model ---
+    if MODEL is None:
+        logger.info(f"Loading InstaNovo (Transformer) model: {TRANSFORMER_MODEL_ID} to {DEVICE}...")
+        try:
+            MODEL, MODEL_CONFIG = InstaNovo.from_pretrained(TRANSFORMER_MODEL_ID)
+            MODEL.to(DEVICE)
+            MODEL.eval()
+            RESIDUE_SET = MODEL.residue_set
+            logger.info("Transformer model loaded successfully.")
+        except Exception as e:
+            logger.error(f"Error loading Transformer model: {e}")
+            raise gr.Error(f"Failed to load InstaNovo Transformer model: {TRANSFORMER_MODEL_ID}. Error: {e}")
+    else:
+         logger.info("Transformer model already loaded.")
 
-def load_model_and_knapsack():
-    """Loads the InstaNovo model and generates/loads the knapsack."""
-    global MODEL, KNAPSACK, MODEL_CONFIG, RESIDUE_SET
-    if MODEL is not None:
-        logger.info("Model already loaded.")
-        return
 
-    logger.info(f"Loading InstaNovo model: {MODEL_ID} to {DEVICE}...")
-    try:
-        MODEL, MODEL_CONFIG = InstaNovo.from_pretrained(MODEL_ID)
-        MODEL.to(DEVICE)
-        MODEL.eval()
-        RESIDUE_SET = MODEL.residue_set
-        logger.info("Model loaded successfully.")
-    except Exception as e:
-        logger.error(f"Error loading model: {e}")
-        raise gr.Error(f"Failed to load InstaNovo model: {MODEL_ID}. Error: {e}")
+    # --- Load Diffusion Model ---
+    if MODEL_PLUS is None:
+        logger.info(f"Loading InstaNovo+ (Diffusion) model: {DIFFUSION_MODEL_ID} to {DEVICE}...")
+        try:
+            MODEL_PLUS, MODEL_PLUS_CONFIG = InstaNovoPlus.from_pretrained(DIFFUSION_MODEL_ID)
+            MODEL_PLUS.to(DEVICE)
+            MODEL_PLUS.eval()
+            if RESIDUE_SET is not None and MODEL_PLUS.residues != RESIDUE_SET:
+                 logger.warning("Residue sets between Transformer and Diffusion models may differ. Using Transformer's set.")
+            elif RESIDUE_SET is None:
+                 RESIDUE_SET = MODEL_PLUS.residues
+
+            logger.info("Diffusion model loaded successfully.")
+        except Exception as e:
+            logger.error(f"Error loading Diffusion model: {e}")
+            gr.Warning(f"Failed to load InstaNovo+ Diffusion model ({DIFFUSION_MODEL_ID}): {e}. Diffusion modes will be unavailable.")
+            MODEL_PLUS = None
+    else:
+        logger.info("Diffusion model already loaded.")
+
 
     # --- Knapsack Handling ---
-    knapsack_exists = (
-        (KNAPSACK_DIR / "parameters.pkl").exists()
-        and (KNAPSACK_DIR / "masses.npy").exists()
-        and (KNAPSACK_DIR / "chart.npy").exists()
-    )
+    # Only attempt knapsack loading/generation if the Transformer model is loaded
+    if MODEL is not None and RESIDUE_SET is not None and KNAPSACK is None:
+        knapsack_exists = (
+            (KNAPSACK_DIR / "parameters.pkl").exists()
+            and (KNAPSACK_DIR / "masses.npy").exists()
+            and (KNAPSACK_DIR / "chart.npy").exists()
+        )
 
-    if knapsack_exists:
-        logger.info(f"Loading pre-generated knapsack from {KNAPSACK_DIR}...")
-        try:
-            KNAPSACK = Knapsack.from_file(str(KNAPSACK_DIR))
-            logger.info("Knapsack loaded successfully.")
-        except Exception as e:
-            logger.info(f"Error loading knapsack: {e}. Will attempt to regenerate.")
-            KNAPSACK = None # Force regeneration
-            knapsack_exists = False # Ensure generation happens
-
-    if not knapsack_exists:
-        logger.info("Knapsack not found or failed to load. Generating knapsack...")
-        if RESIDUE_SET is None:
-            raise gr.Error(
-                "Cannot generate knapsack because ResidueSet failed to load."
-            )
-        try:
-            # Prepare residue masses for knapsack generation (handle negative/zero masses)
-            residue_masses_knapsack = dict(RESIDUE_SET.residue_masses.copy())
-            negative_residues = [
-                k for k, v in residue_masses_knapsack.items() if v <= 0
-            ]
-            if negative_residues:
-                logger.info(f"Warning: Non-positive masses found in residues: {negative_residues}. "
-                      "Excluding from knapsack generation.")
-                for res in negative_residues:
-                    del residue_masses_knapsack[res]
-            # Remove special tokens explicitly if they somehow got mass
-            for special_token in RESIDUE_SET.special_tokens:
-                if special_token in residue_masses_knapsack:
-                    del residue_masses_knapsack[special_token]
-
-            # Ensure residue indices used match those without special/negative masses
-            valid_residue_indices = {
-                res: idx
-                for res, idx in RESIDUE_SET.residue_to_index.items()
-                if res in residue_masses_knapsack
-            }
-
-            KNAPSACK = Knapsack.construct_knapsack(
-                residue_masses=residue_masses_knapsack,
-                residue_indices=valid_residue_indices,  # Use only valid indices
-                max_mass=MAX_MASS,
-                mass_scale=MASS_SCALE,
-            )
-            logger.info(f"Knapsack generated. Saving to {KNAPSACK_DIR}...")
-            KNAPSACK.save(str(KNAPSACK_DIR)) # Save for future runs
-            logger.info("Knapsack saved.")
-        except Exception as e:
-            logger.info(f"Error generating or saving knapsack: {e}")
-            gr.Warning("Failed to generate Knapsack. Knapsack Beam Search will not be available. {e}")
-            KNAPSACK = None # Ensure it's None if generation failed
+        if knapsack_exists:
+            logger.info(f"Loading pre-generated knapsack from {KNAPSACK_DIR}...")
+            try:
+                KNAPSACK = Knapsack.from_file(str(KNAPSACK_DIR))
+                logger.info("Knapsack loaded successfully.")
+            except Exception as e:
+                logger.info(f"Error loading knapsack: {e}. Will attempt to regenerate.")
+                KNAPSACK = None
+                knapsack_exists = False
+
+        if not knapsack_exists:
+            logger.info("Knapsack not found or failed to load. Generating knapsack...")
+            try:
+                residue_masses_knapsack = dict(RESIDUE_SET.residue_masses.copy())
+                special_and_nonpositive = list(RESIDUE_SET.special_tokens) + [
+                    k for k, v in residue_masses_knapsack.items() if v <= 0
+                ]
+                if special_and_nonpositive:
+                     logger.info(f"Excluding special/non-positive mass residues from knapsack: {special_and_nonpositive}")
+                     for res in set(special_and_nonpositive):
+                         if res in residue_masses_knapsack:
+                             del residue_masses_knapsack[res]
+
+                valid_residue_indices = {
+                    res: idx
+                    for res, idx in RESIDUE_SET.residue_to_index.items()
+                    if res in residue_masses_knapsack
+                }
+
+                if not residue_masses_knapsack:
+                    raise ValueError("No valid residues with positive mass found for knapsack generation.")
+
+                KNAPSACK = Knapsack.construct_knapsack(
+                    residue_masses=residue_masses_knapsack,
+                    residue_indices=valid_residue_indices,
+                    max_mass=MAX_MASS,
+                    mass_scale=MASS_SCALE,
+                )
+                logger.info(f"Knapsack generated. Saving to {KNAPSACK_DIR}...")
+                KNAPSACK_DIR.mkdir(parents=True, exist_ok=True)
+                KNAPSACK.save(str(KNAPSACK_DIR))
+                logger.info("Knapsack saved.")
+            except Exception as e:
+                logger.error(f"Error generating or saving knapsack: {e}", exc_info=True)
+                gr.Warning(f"Failed to generate Knapsack. Knapsack Beam Search will not be available. Error: {e}")
+                KNAPSACK = None
+    elif KNAPSACK is not None:
+        logger.info("Knapsack already loaded.")
+    elif MODEL is None:
+         logger.warning("Transformer model not loaded, skipping Knapsack loading/generation.")
 
-# Load the model and knapsack when the script starts
-load_model_and_knapsack()
+
+# Load models and knapsack when the script starts
+load_models_and_knapsack()
 
 
 def create_inference_config(
     input_path: str,
     output_path: str,
-    decoding_method: str,
 ) -> DictConfig:
-    """Creates the OmegaConf DictConfig needed for prediction."""
-    # Load default config if available, otherwise create from scratch
+    """Creates a base OmegaConf DictConfig for prediction environment."""
     if DEFAULT_CONFIG_PATH.exists():
         base_cfg = OmegaConf.load(DEFAULT_CONFIG_PATH)
+        logger.info(f"Loaded base config from {DEFAULT_CONFIG_PATH}")
     else:
          logger.info(f"Warning: Default config not found at {DEFAULT_CONFIG_PATH}. Using minimal config.")
-         # Create a minimal config if default is missing
          base_cfg = OmegaConf.create({
-             "data_path": None,
-             "instanovo_model": MODEL_ID,
-             "output_path": None,
-             "knapsack_path": str(KNAPSACK_DIR),
-             "denovo": True,
-             "refine": False, # Not doing refinement here
-             "num_beams": 1,
-             "max_length": 40,
-             "max_charge": 10,
-             "isotope_error_range": [0, 1],
-             "subset": 1.0,
-             "use_knapsack": False,
-             "save_beams": False,
-             "batch_size": 64, # Adjust as needed
-             "device": DEVICE,
-             "fp16": FP16,
-             "log_interval": 500, # Less relevant for Gradio app
-             "use_basic_logging": True,
-             "filter_precursor_ppm": 20,
-             "filter_confidence": 1e-4,
-             "filter_fdr_threshold": 0.05,
-             "residue_remapping": { # Add default mappings
+             "data_path": None, "instanovo_model": TRANSFORMER_MODEL_ID,
+             "instanovoplus_model": DIFFUSION_MODEL_ID, "output_path": None,
+             "knapsack_path": str(KNAPSACK_DIR), "denovo": True, "refine": True,
+             "num_beams": 1, "max_length": 40, "max_charge": 10,
+             "isotope_error_range": [0, 1], "subset": 1.0, "use_knapsack": False,
+             "save_beams": False, "batch_size": 64, "device": DEVICE, "fp16": FP16,
+             "log_interval": 500, "use_basic_logging": True,
+             "filter_precursor_ppm": 20, "filter_confidence": 1e-4,
+             "filter_fdr_threshold": 0.05, "suppressed_residues": None,
+             "disable_terminal_residues_anywhere": True,
+             "residue_remapping": {
                  "M(ox)": "M[UNIMOD:35]", "M(+15.99)": "M[UNIMOD:35]",
                  "S(p)": "S[UNIMOD:21]", "T(p)": "T[UNIMOD:21]", "Y(p)": "Y[UNIMOD:21]",
                  "S(+79.97)": "S[UNIMOD:21]", "T(+79.97)": "T[UNIMOD:21]", "Y(+79.97)": "Y[UNIMOD:21]",
                  "Q(+0.98)": "Q[UNIMOD:7]", "N(+0.98)": "N[UNIMOD:7]",
                  "Q(+.98)": "Q[UNIMOD:7]", "N(+.98)": "N[UNIMOD:7]",
-                 "C(+57.02)": "C[UNIMOD:4]",
-                 "(+42.01)": "[UNIMOD:1]", "(+43.01)": "[UNIMOD:5]", "(-17.03)": "[UNIMOD:385]",
+                 "C(+57.02)": "C[UNIMOD:4]", "(+42.01)": "[UNIMOD:1]",
+                 "(+43.01)": "[UNIMOD:5]", "(-17.03)": "[UNIMOD:385]",
              },
-             "column_map": { # Add default mappings
+             "column_map": {
                 "Modified sequence": "modified_sequence", "MS/MS m/z": "precursor_mz",
                 "Mass": "precursor_mass", "Charge": "precursor_charge",
                 "Mass values": "mz_array", "Mass spectrum": "mz_array",
@@ -200,256 +234,457 @@ def create_inference_config(
              },
              "index_columns": [
                  "scan_number", "precursor_mz", "precursor_charge",
+                 "retention_time", "spectrum_id", "experiment_name",
              ],
-             # Add other defaults if needed based on errors
          })
 
-    # Override specific parameters
     cfg_overrides = {
-        "data_path": input_path,
-        "output_path": output_path,
-        "device": DEVICE,
-        "fp16": FP16,
-        "denovo": True,
-        "refine": False,
+        "data_path": input_path, "output_path": output_path,
+        "device": DEVICE, "fp16": FP16, "denovo": True,
     }
+    final_cfg = OmegaConf.merge(base_cfg, cfg_overrides)
+    logger.info(f"Created inference config:\n{OmegaConf.to_yaml(final_cfg)}")
+    return final_cfg
 
-    if "Greedy" in decoding_method:
-        cfg_overrides["num_beams"] = 1
-        cfg_overrides["use_knapsack"] = False
-    elif "Knapsack" in decoding_method:
+def _get_transformer_decoder(selection: str, config: DictConfig) -> tuple[Decoder, int, bool]:
+    """Helper to instantiate the correct transformer decoder based on selection."""
+    global MODEL, KNAPSACK
+    if MODEL is None:
+        raise gr.Error("InstaNovo Transformer model not loaded.")
+
+    num_beams = 1
+    use_knapsack = False
+    decoder: Decoder
+
+    if "Greedy" in selection:
+        decoder = GreedyDecoder(
+            model=MODEL,
+            mass_scale=MASS_SCALE,
+            suppressed_residues=config.get("suppressed_residues", None),
+            disable_terminal_residues_anywhere=config.get("disable_terminal_residues_anywhere", True),
+        )
+    elif "Knapsack" in selection:
         if KNAPSACK is None:
-            raise gr.Error(
-                "Knapsack is not available. Cannot use Knapsack Beam Search."
-            )
-        cfg_overrides["num_beams"] = 5
-        cfg_overrides["use_knapsack"] = True
-        cfg_overrides["knapsack_path"] = str(KNAPSACK_DIR)
+            raise gr.Error("Knapsack is not available. Cannot use Knapsack Beam Search.")
+        decoder = KnapsackBeamSearchDecoder(model=MODEL, knapsack=KNAPSACK)
+        num_beams = 5 # Default beam size for knapsack
+        use_knapsack = True
     else:
-        raise ValueError(f"Unknown decoding method: {decoding_method}")
+        raise ValueError(f"Unknown transformer decoder selection: {selection}")
+
+    logger.info(f"Using Transformer decoder: {type(decoder).__name__} (Num beams: {num_beams}, Use Knapsack: {use_knapsack})")
+    return decoder, num_beams, use_knapsack
+
+
+def run_transformer_prediction(dl, config, transformer_decoder_selection):
+    """Runs prediction using only the transformer model."""
+    global RESIDUE_SET
+    if RESIDUE_SET is None:
+        raise gr.Error("ResidueSet not loaded.")
+
+    decoder, num_beams, use_knapsack = _get_transformer_decoder(transformer_decoder_selection, config)
+
+    results_list: list[ScoredSequence | list] = []
+    start_time = time.time()
+    for i, batch in enumerate(dl):
+        spectra, precursors, spectra_mask, _, _ = batch
+        spectra = spectra.to(DEVICE)
+        precursors = precursors.to(DEVICE)
+        spectra_mask = spectra_mask.to(DEVICE)
+
+        with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16):
+            batch_predictions = decoder.decode(
+                spectra=spectra,
+                precursors=precursors,
+                beam_size=num_beams,
+                max_length=config.max_length,
+                mass_tolerance=config.get("filter_precursor_ppm", 20) * 1e-6,
+                max_isotope=config.isotope_error_range[1] if config.isotope_error_range else 1,
+                return_beam=False, # Only top result
+            )
+        results_list.extend(batch_predictions)
+        if (i + 1) % 10 == 0 or (i + 1) == len(dl):
+             logger.info(f"Transformer prediction: Processed batch {i+1}/{len(dl)}")
+
+    end_time = time.time()
+    logger.info(f"Transformer prediction finished in {end_time - start_time:.2f} seconds.")
+    return results_list
+
+def run_diffusion_prediction(dl, config):
+    """Runs prediction using only the diffusion model."""
+    global MODEL_PLUS, RESIDUE_SET
+    if MODEL_PLUS is None or RESIDUE_SET is None:
+        raise gr.Error("InstaNovo+ Diffusion model not loaded.")
+
+    diffusion_decoder = DiffusionDecoder(model=MODEL_PLUS)
+    logger.info(f"Using decoder: {type(diffusion_decoder).__name__}")
+
+    results_sequences = []
+    results_log_probs = []
+    start_time = time.time()
+
+    # Re-create dataloader iterator to get precursor info easily later
+    all_batches = list(dl)
+
+    for i, batch in enumerate(all_batches):
+        spectra, precursors, spectra_mask, _, _ = batch
+        spectra = spectra.to(DEVICE)
+        precursors = precursors.to(DEVICE)
+        spectra_mask = spectra_mask.to(DEVICE)
+
+        with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16):
+            batch_sequences, batch_log_probs = diffusion_decoder.decode(
+                spectra=spectra,
+                spectra_padding_mask=spectra_mask,
+                precursors=precursors,
+                initial_sequence=None,
+            )
+        results_sequences.extend(batch_sequences)
+        results_log_probs.extend(batch_log_probs)
+        if (i + 1) % 10 == 0 or (i + 1) == len(all_batches):
+             logger.info(f"Diffusion prediction: Processed batch {i+1}/{len(all_batches)}")
+
+    end_time = time.time()
+    logger.info(f"Diffusion prediction finished in {end_time - start_time:.2f} seconds.")
+
+    scored_results = []
+    metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range)
+    all_precursors = torch.cat([b[1] for b in all_batches], dim=0) # b[1] is precursors
+
+    for idx, (seq, logp) in enumerate(zip(results_sequences, results_log_probs)):
+         prec_mz = all_precursors[idx, 1].item()
+         prec_ch = int(all_precursors[idx, 2].item())
+         try:
+             _, delta_mass_list = metrics_calc.matches_precursor(seq, prec_mz, prec_ch)
+             min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan")
+         except Exception as e:
+             logger.info(f"Warning: Could not calculate delta mass for diffusion prediction {idx}: {e}")
+             min_abs_ppm = float("nan")
+
+         scored_results.append(
+             ScoredSequence(sequence=seq, mass_error=min_abs_ppm, sequence_log_probability=logp, token_log_probabilities=[])
+         )
+
+    return scored_results
+
+
+def run_refinement_prediction(dl, config, transformer_decoder_selection):
+    """Runs transformer prediction followed by diffusion refinement."""
+    global MODEL, MODEL_PLUS, RESIDUE_SET, MODEL_PLUS_CONFIG
+    if MODEL is None or MODEL_PLUS is None or RESIDUE_SET is None or MODEL_PLUS_CONFIG is None:
+         missing = [m for m, v in [("Transformer", MODEL), ("Diffusion", MODEL_PLUS)] if v is None]
+         raise gr.Error(f"Cannot run refinement: {', '.join(missing)} model not loaded.")
+
+    # 1. Run Transformer Prediction (using selected decoder)
+    logger.info(f"Running Transformer prediction ({transformer_decoder_selection}) for refinement...")
+    transformer_decoder, num_beams, _ = _get_transformer_decoder(transformer_decoder_selection, config) # Get selected decoder
+    transformer_results_list: list[ScoredSequence | list] = []
+
+    all_batches = list(dl) # Store batches
+
+    start_time_transformer = time.time()
+    for i, batch in enumerate(all_batches):
+        spectra, precursors, spectra_mask, _, _ = batch
+        spectra = spectra.to(DEVICE)
+        precursors = precursors.to(DEVICE)
+        spectra_mask = spectra_mask.to(DEVICE)
+
+        with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16):
+            batch_predictions = transformer_decoder.decode(
+                spectra=spectra,
+                precursors=precursors,
+                beam_size=num_beams, # Use selected beam size
+                max_length=config.max_length,
+                mass_tolerance=config.get("filter_precursor_ppm", 20) * 1e-6,
+                max_isotope=config.isotope_error_range[1] if config.isotope_error_range else 1,
+                return_beam=False, # Only top result needed for refinement
+            )
+        transformer_results_list.extend(batch_predictions)
+        if (i + 1) % 10 == 0 or (i + 1) == len(all_batches):
+            logger.info(f"Refinement (Transformer): Processed batch {i+1}/{len(all_batches)}")
+
+    logger.info(f"Transformer prediction for refinement finished in {time.time() - start_time_transformer:.2f} seconds.")
+
+    # 2. Prepare Transformer Predictions as Initial Sequences for Diffusion
+    logger.info("Encoding transformer predictions for diffusion input...")
+    encoded_transformer_preds = []
+    max_len_diffusion = MODEL_PLUS_CONFIG.get("max_length", 40)
+
+    for res in transformer_results_list:
+        if isinstance(res, ScoredSequence) and res.sequence:
+             # Encode sequence *without* EOS for diffusion input.
+             encoded = RESIDUE_SET.encode(res.sequence, add_eos=False, return_tensor='pt')
+        else:
+            # If transformer failed, provide a dummy PAD sequence
+             encoded = torch.full((max_len_diffusion,), RESIDUE_SET.PAD_INDEX, dtype=torch.long)
+
+
+        # Pad or truncate to the diffusion model's max length
+        current_len = encoded.shape[0]
+        if current_len > max_len_diffusion:
+             logger.warning(f"Transformer prediction exceeded diffusion max length ({max_len_diffusion}). Truncating.")
+             encoded = encoded[:max_len_diffusion]
+        elif current_len < max_len_diffusion:
+            padding = torch.full((max_len_diffusion - current_len,), RESIDUE_SET.PAD_INDEX, dtype=torch.long)
+            encoded = torch.cat((encoded, padding))
+
+        encoded_transformer_preds.append(encoded)
+
+    if not encoded_transformer_preds:
+        raise gr.Error("Transformer prediction yielded no results to refine.")
+    encoded_transformer_preds_tensor = torch.stack(encoded_transformer_preds).to(DEVICE)
+    logger.info(f"Encoded {encoded_transformer_preds_tensor.shape[0]} sequences for diffusion.")
+
+
+    # 3. Run Diffusion Refinement
+    logger.info("Running Diffusion refinement...")
+    diffusion_decoder = DiffusionDecoder(model=MODEL_PLUS)
+    refined_sequences = []
+    refined_log_probs = []
+    start_time_diffusion = time.time()
+
+    current_idx = 0
+    for i, batch in enumerate(all_batches):
+        spectra, precursors, spectra_mask, _, _ = batch
+        spectra = spectra.to(DEVICE)
+        precursors = precursors.to(DEVICE)
+        spectra_mask = spectra_mask.to(DEVICE)
+
+        batch_size = spectra.shape[0]
+        initial_sequences_batch = encoded_transformer_preds_tensor[current_idx : current_idx + batch_size]
+        current_idx += batch_size
+
+        if initial_sequences_batch.shape[0] != batch_size:
+             logger.error(f"Batch size mismatch during refinement: expected {batch_size}, got {initial_sequences_batch.shape[0]}")
+             continue # Skip batch?
+
+        with torch.no_grad(), torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16):
+            batch_refined_seqs, batch_refined_logp = diffusion_decoder.decode(
+                spectra=spectra,
+                spectra_padding_mask=spectra_mask,
+                precursors=precursors,
+                initial_sequence=initial_sequences_batch,
+                start_step=DIFFUSION_START_STEP,
+            )
+        refined_sequences.extend(batch_refined_seqs)
+        refined_log_probs.extend(batch_refined_logp)
+        if (i + 1) % 10 == 0 or (i + 1) == len(all_batches):
+             logger.info(f"Refinement (Diffusion): Processed batch {i+1}/{len(all_batches)}")
+
+    logger.info(f"Diffusion refinement finished in {time.time() - start_time_diffusion:.2f} seconds.")
+
+    # 4. Combine and Format Results
+    all_precursors = torch.cat([b[1] for b in all_batches], dim=0) # b[1] is precursors
+    metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range)
+    combined_results = []
+    for idx, (transformer_res, refined_seq, refined_logp) in enumerate(zip(transformer_results_list, refined_sequences, refined_log_probs)):
+         prec_mz = all_precursors[idx, 1].item()
+         prec_ch = int(all_precursors[idx, 2].item())
+         try:
+             _, delta_mass_list = metrics_calc.matches_precursor(refined_seq, prec_mz, prec_ch)
+             min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan")
+         except Exception as e:
+             logger.info(f"Warning: Could not calculate delta mass for refined prediction {idx}: {e}")
+             min_abs_ppm = float("nan")
+
+         combined_data = {
+             "transformer_prediction": "".join(transformer_res.sequence) if isinstance(transformer_res, ScoredSequence) else "",
+             "transformer_log_probability": transformer_res.sequence_log_probability if isinstance(transformer_res, ScoredSequence) else float('-inf'),
+             "refined_prediction": "".join(refined_seq),
+             "refined_log_probability": refined_logp,
+             "refined_delta_mass_ppm": min_abs_ppm,
+         }
+         combined_results.append(combined_data)
+
+    return combined_results
 
-    # Merge base config with overrides
-    final_cfg = OmegaConf.merge(base_cfg, cfg_overrides)
-    return final_cfg
 
 @spaces.GPU
-def predict_peptides(input_file, decoding_method):
+def predict_peptides(input_file, mode_selection, transformer_decoder_selection):
     """
-    Main function to load data, run prediction, and return results.
+    Main function to load data, select mode, run prediction, and return results.
     """
-    if MODEL is None or RESIDUE_SET is None or MODEL_CONFIG is None:
-        load_model_and_knapsack()  # Attempt to reload if None (e.g., after space restart)
+    # Ensure models are loaded
+    if MODEL is None or RESIDUE_SET is None:
+        load_models_and_knapsack() # Try reload
         if MODEL is None:
-            raise gr.Error("InstaNovo model is not loaded. Cannot perform prediction.")
+            raise gr.Error("InstaNovo Transformer model failed to load. Cannot perform prediction.")
+    if ("Refinement" in mode_selection or "InstaNovo+" in mode_selection) and MODEL_PLUS is None:
+         load_models_and_knapsack() # Try reload diffusion
+         if MODEL_PLUS is None:
+             raise gr.Error("InstaNovo+ Diffusion model failed to load. Cannot perform Refinement or InstaNovo+ Only prediction.")
+    if "Knapsack" in transformer_decoder_selection and KNAPSACK is None:
+         load_models_and_knapsack() # Try reload knapsack
+         if KNAPSACK is None:
+             raise gr.Error("Knapsack failed to load. Cannot use Knapsack Beam Search.")
+
 
     if input_file is None:
         raise gr.Error("Please upload a mass spectrometry file.")
 
-    input_path = input_file.name # Gradio provides the path in .name
-    logger.info(f"Processing file: {input_path}")
-    logger.info(f"Using decoding method: {decoding_method}")
+    input_path = input_file.name
+    logger.info(f"--- New Prediction Request ---")
+    logger.info(f"Input File: {input_path}")
+    logger.info(f"Selected Mode: {mode_selection}")
+    if "Refinement" in mode_selection or "InstaNovo Only" in mode_selection:
+        logger.info(f"Selected Transformer Decoder: {transformer_decoder_selection}")
 
-    # Create a temporary file for the output CSV
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as temp_out:
-        output_csv_path = temp_out.name
+    # Create temp output file
+    gradio_tmp_dir = os.environ.get("GRADIO_TEMP_DIR", "/tmp")
+    try:
+        with tempfile.NamedTemporaryFile(dir=gradio_tmp_dir, delete=False, suffix=".csv") as temp_out:
+            output_csv_path = temp_out.name
+        logger.info(f"Temporary output path: {output_csv_path}")
+    except Exception as e:
+         logger.error(f"Failed to create temporary file in {gradio_tmp_dir}: {e}")
+         raise gr.Error(f"Failed to create temporary output file: {e}")
 
     try:
-        # 1. Create Config
-        config = create_inference_config(input_path, output_csv_path, decoding_method)
-        logger.info(f"Inference Config:\n{OmegaConf.to_yaml(config)}")
+        config = create_inference_config(input_path, output_csv_path)
 
-        # 2. Load Data using SpectrumDataFrame
         logger.info("Loading spectrum data...")
         try:
+            # Load data eagerly
             sdf = SpectrumDataFrame.load(
-                config.data_path,
-                lazy=False,  # Load eagerly for Gradio simplicity
-                is_annotated=False,  # De novo mode
-                column_mapping=config.get("column_map", None),
-                shuffle=False,
-                verbose=True,  # Print loading logs
+                config.data_path, lazy=False, is_annotated=False,
+                column_mapping=config.get("column_map", None), shuffle=False, verbose=True,
             )
-            # Apply charge filter like in CLI
             original_size = len(sdf)
             max_charge = config.get("max_charge", 10)
-            sdf.filter_rows(
-                lambda row: (row["precursor_charge"] <= max_charge)
-                and (row["precursor_charge"] > 0)
-            )
-            if len(sdf) < original_size:
-                logger.info(f"Warning: Filtered {original_size - len(sdf)} spectra with charge > {max_charge} or <= 0.")
+            if "precursor_charge" in sdf.df.columns:
+                sdf.filter_rows(
+                    lambda row: ("precursor_charge" in row and row["precursor_charge"] is not None and 0 < row["precursor_charge"] <= max_charge)
+                )
+                if len(sdf) < original_size:
+                    logger.info(f"Warning: Filtered {original_size - len(sdf)} spectra with invalid or out-of-range charge (<=0 or >{max_charge}).")
+            else:
+                 logger.warning("Column 'precursor_charge' not found. Cannot filter by charge.")
 
             if len(sdf) == 0:
                  raise gr.Error("No valid spectra found in the uploaded file after filtering.")
             logger.info(f"Data loaded: {len(sdf)} spectra.")
+            index_cols_present = [col for col in config.index_columns if col in sdf.df.columns]
+            base_df_pd = sdf.df.select(index_cols_present).to_pandas()
+
         except Exception as e:
-            logger.info(f"Error loading data: {e}")
+            logger.error(f"Error loading data: {e}", exc_info=True)
             raise gr.Error(f"Failed to load or process the spectrum file. Error: {e}")
 
-        # 3. Prepare Dataset and DataLoader
+        if RESIDUE_SET is None: raise gr.Error("Residue set not loaded.") # Should not happen if model loaded
+
+        # --- Prepare DataLoader ---
+        # Use reverse_peptide=True for Transformer steps, False for Diffusion-only
+        reverse_for_transformer = "InstaNovo+ Only" not in mode_selection
         ds = SpectrumDataset(
-            sdf,
-            RESIDUE_SET,
-            MODEL_CONFIG.get("n_peaks", 200),
-            return_str=True,  # Needed for greedy/beam search targets later (though not used here)
-            annotated=False,
-            pad_spectrum_max_length=config.get("compile_model", False)
-            or config.get("use_flash_attention", False),
+            sdf, RESIDUE_SET,
+            MODEL_CONFIG.get("n_peaks", 200) if MODEL_CONFIG else 200,
+            return_str=True, annotated=False,
+            pad_spectrum_max_length=config.get("compile_model", False) or config.get("use_flash_attention", False),
             bin_spectra=config.get("conv_peak_encoder", False),
+            peptide_pad_length=config.get("max_length", 40) if config.get("compile_model", False) else 0,
+            reverse_peptide=reverse_for_transformer, # Key change based on mode
+            diffusion="InstaNovo+ Only" in mode_selection # Signal if input is for diffusion
         )
-        dl = DataLoader(
-            ds,
-            batch_size=config.batch_size,
-            num_workers=0,  # Required by SpectrumDataFrame
-            shuffle=False,  # Required by SpectrumDataFrame
-            collate_fn=collate_batch,
-        )
+        dl = DataLoader(ds, batch_size=config.batch_size, num_workers=0, shuffle=False, collate_fn=collate_batch)
+
+        # --- Run Prediction ---
+        results_data = None
+        output_headers = index_cols_present[:]
+
+        if "InstaNovo Only" in mode_selection:
+             output_headers.extend(["prediction", "log_probability", "delta_mass_ppm", "token_log_probabilities"])
+             transformer_results = run_transformer_prediction(dl, config, transformer_decoder_selection)
+             results_data = []
+             metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range)
+             for i, res in enumerate(transformer_results):
+                 row_data = {}
+                 if isinstance(res, ScoredSequence) and res.sequence:
+                     row_data["prediction"] = "".join(res.sequence)
+                     row_data["log_probability"] = f"{res.sequence_log_probability:.4f}"
+                     row_data["token_log_probabilities"] = ", ".join(f"{p:.4f}" for p in res.token_log_probabilities)
+                     try:
+                         prec_mz = base_df_pd.loc[i, "precursor_mz"]
+                         prec_ch = base_df_pd.loc[i, "precursor_charge"]
+                         _, delta_mass_list = metrics_calc.matches_precursor(res.sequence, prec_mz, prec_ch)
+                         min_abs_ppm = min(abs(p) for p in delta_mass_list) if delta_mass_list else float("nan")
+                         row_data["delta_mass_ppm"] = f"{min_abs_ppm:.2f}"
+                     except Exception as e:
+                          logger.warning(f"Could not calculate delta mass for Tx prediction {i}: {e}")
+                          row_data["delta_mass_ppm"] = "N/A"
+                 else:
+                     row_data.update({k: "N/A" for k in ["prediction", "log_probability", "delta_mass_ppm", "token_log_probabilities"]})
+                     row_data["prediction"] = "" # Ensure empty string for failed preds
+                     row_data["token_log_probabilities"] = ""
+                 results_data.append(row_data)
+
+        elif "InstaNovo+ Only" in mode_selection:
+             output_headers.extend(["prediction", "log_probability", "delta_mass_ppm"])
+             diffusion_results = run_diffusion_prediction(dl, config)
+             results_data = []
+             for res in diffusion_results:
+                 row_data = {}
+                 if isinstance(res, ScoredSequence) and res.sequence:
+                     row_data["prediction"] = "".join(res.sequence)
+                     row_data["log_probability"] = f"{res.sequence_log_probability:.4f}" # Avg loss
+                     row_data["delta_mass_ppm"] = f"{res.mass_error:.2f}" if not np.isnan(res.mass_error) else "N/A" # ppm
+                 else:
+                     row_data.update({k: "N/A" for k in ["prediction", "log_probability", "delta_mass_ppm"]})
+                     row_data["prediction"] = ""
+                 results_data.append(row_data)
+
+        elif "Refinement" in mode_selection:
+             output_headers.extend([
+                 "transformer_prediction", "transformer_log_probability",
+                 "refined_prediction", "refined_log_probability", "refined_delta_mass_ppm"
+             ])
+             # Pass the selected transformer decoder to the refinement function
+             results_data = run_refinement_prediction(dl, config, transformer_decoder_selection)
+             for row in results_data:
+                 # Format numbers after getting the list of dicts
+                 row["transformer_log_probability"] = f"{row['transformer_log_probability']:.4f}" if isinstance(row['transformer_log_probability'], (float, int)) else "N/A"
+                 row["refined_log_probability"] = f"{row['refined_log_probability']:.4f}" if isinstance(row['refined_log_probability'], (float, int)) else "N/A"
+                 row["refined_delta_mass_ppm"] = f"{row['refined_delta_mass_ppm']:.2f}" if isinstance(row['refined_delta_mass_ppm'], (float, int)) and not np.isnan(row['refined_delta_mass_ppm']) else "N/A"
+
 
-        # 4. Select Decoder
-        logger.info("Initializing decoder...")
-        decoder: Decoder
-        if config.use_knapsack:
-            if KNAPSACK is None:
-                # This check should ideally be earlier, but double-check
-                raise gr.Error(
-                    "Knapsack is required for Knapsack Beam Search but is not available."
-                )
-            # KnapsackBeamSearchDecoder doesn't directly load from path in this version?
-            # We load Knapsack globally, so just pass it.
-            # If it needed path: decoder = KnapsackBeamSearchDecoder.from_file(model=MODEL, path=config.knapsack_path)
-            decoder = KnapsackBeamSearchDecoder(model=MODEL, knapsack=KNAPSACK)
-        elif config.num_beams > 1:
-             # BeamSearchDecoder is available but not explicitly requested, use Greedy for num_beams=1
-             logger.info(f"Warning: num_beams={config.num_beams} > 1 but only Greedy and Knapsack Beam Search are implemented in this app. Defaulting to Greedy.")
-             decoder = GreedyDecoder(model=MODEL, mass_scale=MASS_SCALE)
         else:
-             decoder = GreedyDecoder(
-                 model=MODEL,
-                 mass_scale=MASS_SCALE,
-                 # Add suppression options if needed from config
-                 suppressed_residues=config.get("suppressed_residues", None),
-                 disable_terminal_residues_anywhere=config.get("disable_terminal_residues_anywhere", True),
-             )
-        logger.info(f"Using decoder: {type(decoder).__name__}")
-
-        # 5. Run Prediction Loop (Adapted from instanovo/transformer/predict.py)
-        logger.info("Starting prediction...")
-        start_time = time.time()
-        results_list: list[
-            ScoredSequence | list
-        ] = []  # Store ScoredSequence or empty list
-
-        for i, batch in enumerate(dl):
-            spectra, precursors, spectra_mask, _, _ = (
-                batch  # Ignore peptides/masks for de novo
-            )
-            spectra = spectra.to(DEVICE)
-            precursors = precursors.to(DEVICE)
-            spectra_mask = spectra_mask.to(DEVICE)
-
-            with (
-                torch.no_grad(),
-                torch.amp.autocast(DEVICE, dtype=torch.float16, enabled=FP16),
-            ):
-                # Beam search decoder might return list[list[ScoredSequence]] if return_beam=True
-                # Greedy decoder returns list[ScoredSequence]
-                # KnapsackBeamSearchDecoder returns list[ScoredSequence] or list[list[ScoredSequence]]
-                batch_predictions = decoder.decode(
-                    spectra=spectra,
-                    precursors=precursors,
-                    beam_size=config.num_beams,
-                    max_length=config.max_length,
-                    # Knapsack/Beam Search specific params if needed
-                    mass_tolerance=config.get("filter_precursor_ppm", 20)
-                    * 1e-6,  # Convert ppm to relative
-                    max_isotope=config.isotope_error_range[1]
-                    if config.isotope_error_range
-                    else 1,
-                    return_beam=False,  # Only get the top prediction for simplicity
-                )
-            results_list.extend(batch_predictions) # Should be list[ScoredSequence] or list[list]
-            logger.info(f"Processed batch {i+1}/{len(dl)}")
-
-        end_time = time.time()
-        logger.info(f"Prediction finished in {end_time - start_time:.2f} seconds.")
-
-        # 6. Format Results
-        logger.info("Formatting results...")
-        output_data = []
-        # Use sdf index columns + prediction results
-        index_cols = [col for col in config.index_columns if col in sdf.df.columns]
-        base_df_pd = sdf.df.select(index_cols).to_pandas()  # Get base info
-
-        metrics_calc = Metrics(RESIDUE_SET, config.isotope_error_range)
-
-        for i, res in enumerate(results_list):
-            row_data = base_df_pd.iloc[i].to_dict()  # Get corresponding input data
-            if isinstance(res, ScoredSequence) and res.sequence:
-                sequence_str = "".join(res.sequence)
-                row_data["prediction"] = sequence_str
-                row_data["log_probability"] = f"{res.sequence_log_probability:.4f}"
-                # Use metrics to calculate delta mass ppm for the top prediction
-                try:
-                    _, delta_mass_list = metrics_calc.matches_precursor(
-                        res.sequence,
-                        row_data["precursor_mz"],
-                        row_data["precursor_charge"],
-                    )
-                    # Find the smallest absolute ppm error across isotopes
-                    min_abs_ppm = (
-                        min(abs(p) for p in delta_mass_list)
-                        if delta_mass_list
-                        else float("nan")
-                    )
-                    row_data["delta_mass_ppm"] = f"{min_abs_ppm:.2f}"
-                except Exception as e:
-                     logger.info(f"Warning: Could not calculate delta mass for prediction {i}: {e}")
-                     row_data["delta_mass_ppm"] = "N/A"
+            raise ValueError(f"Unknown mode selection: {mode_selection}")
 
-            else:
-                row_data["prediction"] = ""
-                row_data["log_probability"] = "N/A"
-                row_data["delta_mass_ppm"] = "N/A"
-            output_data.append(row_data)
-
-        output_df = pl.DataFrame(output_data)
-
-        # Ensure specific columns are present and ordered
-        display_cols = [
-            "scan_number",
-            "precursor_mz",
-            "precursor_charge",
-            "prediction",
-            "log_probability",
-            "delta_mass_ppm",
-        ]
-        final_display_cols = []
-        for col in display_cols:
-            if col in output_df.columns:
-                final_display_cols.append(col)
-            else:
-                 logger.info(f"Warning: Expected display column '{col}' not found in results.")
 
-        # Add any remaining index columns that weren't in display_cols
-        for col in index_cols:
-            if col not in final_display_cols and col in output_df.columns:
-                final_display_cols.append(col)
+        # --- Combine, Save, Return ---
+        logger.info("Combining results...")
+        if results_data is None: raise gr.Error("Prediction did not produce results.")
+
+        results_df = pl.DataFrame(results_data)
+        # Ensure base_df_pd has unique index if using join, or just concat horizontally if order is guaranteed
+        base_df_pl = pl.from_pandas(base_df_pd.reset_index(drop=True))
 
-        output_df_display = output_df.select(final_display_cols)
+        # Simple horizontal concat assuming order is preserved by dataloader (shuffle=False)
+        if len(base_df_pl) == len(results_df):
+            final_df = pl.concat([base_df_pl, results_df], how="horizontal")
+        else:
+            logger.error(f"Length mismatch between base data ({len(base_df_pl)}) and results ({len(results_df)}). Cannot reliably combine.")
+            # Fallback or error? Let's just use results for now, but log error.
+            final_df = results_df # Display only results in case of mismatch
+
+        logger.info(f"Saving full results to {output_csv_path}...")
+        final_df.write_csv(output_csv_path)
+        logger.info("Save complete.")
 
-        # 7. Save full results to CSV
-        logger.info(f"Saving results to {output_csv_path}...")
-        output_df.write_csv(output_csv_path)
+        # Select display columns - make sure they exist in final_df
+        display_cols_final = [col for col in output_headers if col in final_df.columns]
+        display_df = final_df.select(display_cols_final)
 
-        # Return DataFrame for display and path for download
-        return output_df_display.to_pandas(), output_csv_path
+        logger.info("--- Prediction Request Complete ---")
+        return display_df.to_pandas(), output_csv_path
 
     except Exception as e:
-        logger.info(f"An error occurred during prediction: {e}")
-        # Clean up the temporary output file if it exists
-        if os.path.exists(output_csv_path):
-            os.remove(output_csv_path)
-        # Re-raise as Gradio error
+        logger.error(f"An error occurred during prediction: {e}", exc_info=True)
+        if 'output_csv_path' in locals() and os.path.exists(output_csv_path):
+            try:
+                os.remove(output_csv_path)
+                logger.info(f"Removed temporary file {output_csv_path}")
+            except OSError:
+                 logger.error(f"Failed to remove temporary file {output_csv_path}")
         raise gr.Error(f"Prediction failed: {e}")
 
 
@@ -458,29 +693,29 @@ css = """
 .gradio-container { font-family: sans-serif; }
 .gr-button { color: white; border-color: black; background: black; }
 footer { display: none !important; }
-/* Optional: Add some margin below the logo */
 .logo-container img { margin-bottom: 1rem; }
+.feedback { font-size: 0.9rem; color: gray; }
 """
 
 with gr.Blocks(
     css=css, theme=gr.themes.Default(primary_hue="blue", secondary_hue="blue")
 ) as demo:
-    # --- Logo Display ---
     gr.Markdown(
         """
         <div style="text-align: center;" class="logo-container">
           <img src='/gradio_api/file=assets/instanovo.svg' alt="InstaNovo Logo" width="300" style="display: block; margin: 0 auto;">
         </div>
         """,
-        elem_classes="logo-container",  # Optional class for CSS targeting
+        elem_classes="logo-container",
     )
 
-    # --- App Content ---
     gr.Markdown(
-        """
-        # 🚀 _De Novo_ Peptide Sequencing with InstaNovo 
-        Upload your mass spectrometry data file (.mgf, .mzml, or .mzxml) and get peptide sequence predictions using InstaNovo.
-        Choose between fast Greedy Search or more accurate but slower Knapsack Beam Search.
+        f"""
+        # 🚀 _De Novo_ Peptide Sequencing with InstaNovo
+        Upload your mass spectrometry data file (.mgf, .mzml, or .mzxml) and get peptide sequence predictions.
+        Choose your prediction method and decoding options.
+
+        **Note:** The InstaNovo+ model `{DIFFUSION_MODEL_ID}` is an alpha release.
         """
     )
     with gr.Row():
@@ -489,73 +724,114 @@ with gr.Blocks(
                 label="Upload Mass Spectrometry File (.mgf, .mzml, .mzxml)",
                 file_types=[".mgf", ".mzml", ".mzxml"],
             )
-            decoding_method = gr.Radio(
+            mode_selection = gr.Radio(
                 [
-                    "Greedy Search (Fast, resonably accurate)",
-                    "Knapsack Beam Search (More accurate, but slower)",
+                    "InstaNovo + Refinement (Default, Recommended)",
+                    "InstaNovo Only (Transformer)",
+                    "InstaNovo+ Only (Diffusion, Alpha)",
                 ],
-                label="Decoding Method",
-                value="Greedy Search (Fast, resonably accurate)",  # Default to fast method
+                label="Prediction Mode",
+                value="InstaNovo + Refinement (Default, Recommended)",
             )
+            # Transformer decoder selection - visible for relevant modes
+            transformer_decoder_selection = gr.Radio(
+                [
+                    "Greedy Search (Fast)",
+                    # Knapsack option added dynamically based on KNAPSACK availability
+                ],
+                label="Transformer Decoding Method",
+                value="Greedy Search (Fast)",
+                visible=True, # Start visible as default mode uses it
+                interactive=True,
+            )
+
             submit_btn = gr.Button("Predict Sequences", variant="primary")
+
+            # --- Control Visibility & Choices ---
+            def update_transformer_options(mode):
+                # Show decoder selection if mode uses the transformer
+                show_decoder = "InstaNovo+ Only" not in mode
+                # Update choices based on knapsack availability
+                knapsack_available = KNAPSACK is not None
+                choices = ["Greedy Search (Fast)"]
+                if knapsack_available:
+                    choices.append("Knapsack Beam Search (Accurate, Slower)")
+                else:
+                    logger.info("Knapsack check: Not available, disabling Knapsack Beam Search option.")
+                # Reset to Greedy if Knapsack was selected but becomes unavailable
+                current_value = "Greedy Search (Fast)" # Default reset value
+                return gr.update(visible=show_decoder, choices=choices, value=current_value)
+
+            mode_selection.change(
+                fn=update_transformer_options,
+                inputs=mode_selection,
+                outputs=transformer_decoder_selection,
+            )
+            # Initial check in case knapsack fails on startup
+            # This requires JS or a different approach in Gradio.
+            # For simplicity, we rely on the check during prediction.
+            # We can set initial choices based on load status here though.
+            initial_choices = ["Greedy Search (Fast)"]
+            if KNAPSACK is not None:
+                initial_choices.append("Knapsack Beam Search (Accurate, Slower)")
+            transformer_decoder_selection.choices = initial_choices
+
+
         with gr.Column(scale=2):
             output_df = gr.DataFrame(
-                label="Prediction Results",
-                headers=[
-                    "scan_number",
-                    "precursor_mz",
-                    "precursor_charge",
-                    "prediction",
-                    "log_probability",
-                    "delta_mass_ppm",
-                ],
-                wrap=True,
+                label="Prediction Results Preview",
+                headers=["scan_number", "prediction", "log_probability", "delta_mass_ppm"]
             )
             output_file = gr.File(label="Download Full Results (CSV)")
 
     submit_btn.click(
         predict_peptides,
-        inputs=[input_file, decoding_method],
+        inputs=[input_file, mode_selection, transformer_decoder_selection],
         outputs=[output_df, output_file],
     )
 
     gr.Examples(
         [
-            ["assets/sample_spectra.mgf", "Greedy Search (Fast, resonably accurate)"],
-            [
-                "assets/sample_spectra.mgf",
-                "Knapsack Beam Search (More accurate, but slower)",
-            ],
+            ["assets/sample_spectra.mgf", "InstaNovo + Refinement (Default, Recommended)", "Greedy Search (Fast)"],
+            ["assets/sample_spectra.mgf", "InstaNovo + Refinement (Default, Recommended)", "Knapsack Beam Search (Accurate, Slower)"],
+            ["assets/sample_spectra.mgf", "InstaNovo Only (Transformer)", "Greedy Search (Fast)"],
+            ["assets/sample_spectra.mgf", "InstaNovo Only (Transformer)", "Knapsack Beam Search (Accurate, Slower)"],
+            ["assets/sample_spectra.mgf", "InstaNovo+ Only (Diffusion, Alpha)", "Greedy Search (Fast)"],
         ],
-        inputs=[input_file, decoding_method],
-        outputs=[output_df, output_file],
-        fn=predict_peptides,
-        cache_examples=False,  # Re-run examples if needed
-        label="Example Usage",
+        inputs=[input_file, mode_selection, transformer_decoder_selection],
+        # outputs=[output_df, output_file],
+        cache_examples=False,
+        label="Example Usage (Note: Knapsack examples require Knapsack to be available)",
     )
 
     gr.Markdown(
-        """
+        f"""
          **Notes:**
-         *   Predictions are based on the [InstaNovo](https://github.com/instadeepai/InstaNovo) model `{MODEL_ID}`.
-         *   Knapsack Beam Search uses pre-calculated mass constraints and yields better results but takes longer.
-         *   `delta_mass_ppm` shows the lowest absolute precursor mass error (in ppm) across potential isotopes (0-1 neutron).
-         *   Ensure your input file format is correctly specified. Large files may take time to process.
-         """.format(MODEL_ID=MODEL_ID)
+         *   Predictions use `{TRANSFORMER_MODEL_ID}` (Transformer) and `{DIFFUSION_MODEL_ID}` (Diffusion, Alpha).
+         *   **Refinement Mode:** Runs initial prediction with the selected Transformer method (Greedy/Knapsack), then refines using InstaNovo+.
+         *   **InstaNovo Only Mode:** Uses only the Transformer with the selected decoding method.
+         *   **InstaNovo+ Only Mode:** Predicts directly using the Diffusion model (alpha version).
+         *   `delta_mass_ppm` shows the lowest absolute precursor mass error (ppm) across isotopes 0-1 for the final sequence.
+         *   Knapsack Beam Search requires a pre-computed knapsack file. If unavailable, the option will be disabled.
+         *   Check logs for progress, especially for large files or slower methods.
+         """,
+         elem_classes="feedback"
     )
 
-    # Add logging component 
-    with gr.Accordion("Application Logs", open=True):
+    with gr.Accordion("Application Logs", open=False):
         log_display = Log(log_file, dark=True, height=300)
-        
-    gr.Textbox(
+
+    gr.Markdown(
         value="""
+If you use InstaNovo in your research, please cite:
+
+```bibtex
 @article{eloff_kalogeropoulos_2025_instanovo,
 	title        = {InstaNovo enables diffusion-powered de novo peptide sequencing in large-scale proteomics experiments},
-	author       = {Kevin Eloff and Konstantinos Kalogeropoulos and Amandla Mabona and Oliver Morell and Rachel Catzel and 
-                    Esperanza Rivera-de-Torre and Jakob Berg Jespersen and Wesley Williams and Sam P. B. van Beljouw and 
-                    Marcin J. Skwark and Andreas Hougaard Laustsen and Stan J. J. Brouns and Anne Ljungars and Erwin M. 
-                    Schoof and Jeroen Van Goey and Ulrich auf dem Keller and Karim Beguir and Nicolas Lopez Carranza and 
+	author       = {Kevin Eloff and Konstantinos Kalogeropoulos and Amandla Mabona and Oliver Morell and Rachel Catzel and
+                    Esperanza Rivera-de-Torre and Jakob Berg Jespersen and Wesley Williams and Sam P. B. van Beljouw and
+                    Marcin J. Skwark and Andreas Hougaard Laustsen and Stan J. J. Brouns and Anne Ljungars and Erwin M.
+                    Schoof and Jeroen Van Goey and Ulrich auf dem Keller and Karim Beguir and Nicolas Lopez Carranza and
                     Timothy P. Jenkins},
 	year         = 2025,
 	month        = {Mar},
@@ -566,8 +842,7 @@ with gr.Blocks(
 }
 """,
         show_copy_button=True,
-        label="If you use InstaNovo in your research, please cite:",
-        interactive=False,
+        label="If you use InstaNovo in your research, please cite:"
     )
 
 # --- Launch the App ---
@@ -576,4 +851,5 @@ if __name__ == "__main__":
     # Set server_name="0.0.0.0" to allow access from network if needed
     # demo.launch(server_name="0.0.0.0", server_port=7860)
     # For Hugging Face Spaces, just demo.launch() is usually sufficient
-    demo.launch(share=True)  # For local testing with public URL
+    demo.launch()
+    # demo.launch(share=True)  # For local testing with public URL