app.py

import os
import re
import shutil
import subprocess
import sys
import threading
from itertools import groupby
from pathlib import Path
from tempfile import NamedTemporaryFile

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import spaces
import torch
from transformers import T5EncoderModel, T5Tokenizer

# -----------------------------------------------------------------------------
# Reproducible upstream sources: these match the working TREAD Colab demo.
# -----------------------------------------------------------------------------
TREAD_REPO = "https://github.com/KYQiu21/TREAD.git"
TREAD_COMMIT = "7a7b3f571778cb89035c798c88df7eacb63ed2c1"
PROTT5_MODEL = "Rostlab/prot_t5_xl_half_uniref50-enc"

APP_DIR = Path(__file__).resolve().parent
TREAD_DIR = APP_DIR / ".tread_source"

# Web-demo guardrails only; these are not biological/model validity limits.
MAX_GPU_SEQUENCE_LENGTH = int(os.getenv("TREAD_MAX_GPU_SEQUENCE_LENGTH", "1200"))
MAX_CPU_SEQUENCE_LENGTH = int(os.getenv("TREAD_MAX_CPU_SEQUENCE_LENGTH", "600"))

REPEAT_HEADS = [
    "Any repeat (segmentation head)",
    "Alpha solenoid",
    "TIM-barrel",
    "Beta-propeller",
    "Beta-barrel",
    "Beta-solenoid",
    "Alpha/beta solenoid",
]
TYPE_HEAD_INDEX = {
    "Alpha solenoid": 0,
    "TIM-barrel": 1,
    "Beta-propeller": 2,
    "Beta-barrel": 3,
    "Beta-solenoid": 4,
    "Alpha/beta solenoid": 5,
}

EXAMPLE_REPEAT = (
    "MMKRNILAVIVPALLVAGTANAAEIYNKDGNKVDLYGKAVGLHYFSKGNGENSYGGNGDMTYARLGFKGETQINSDLTGYGQWEYNFQGNNSEGADAQTGNKTRLAFAGLKYADVGSFDYGRNYGVVYDALGYTDMLPEFGGDTAYSDDFFVGRVGGVATYRNSNFFGLVDGLNFAVQYLGKNERDTARRSNGDGVGGSISYEYEGFGIVGAYGAADRTNLQEAQPLGNGKKAEQWATGLKYDANNIYLAANYGETRNATPITNKFTNTSGFANKTQDVLLVAQYQFDFGLRPSIAYTKSKAKDVEGIGDVDLVNYFEVGATYYFNKNMSTYVDYIINQIDSDNKLGVGSDDTVAVGIVYQF"
)
EXAMPLE_PROPELLER = (
    "MEQVLYLGSYTKRESKGVHQIILDTDKKELRDYRLIAEVDSPTYLDLSADKGTLYSISKTDEGGGITSFKKNENGTYDKVAEISAEGSAPCYIHYDEDKKLIFTANYHGGYLTVYKENADGSFTMSDRAQHEGSSIHENQTIPHVHYSALSPDKKFLLACDLGTDEVYTYTVSDEGKLTEAARYKATPGTGPRHLVFHPNGKVAYLFGELSSDVEVLAYEAATGTFSLLQVITTIPAEHTGFNGGAAIRISADGKFVYASNRGHDSLVVYAVSEDGETLSLVEYVPTEGNTPRDFNLDPSGQFVIVAHQDSDNLTLFERDATTGKLTLVQKDVYAPECVCVFY"
)

# -----------------------------------------------------------------------------
# Model state
#
# ZeroGPU path:
#   The GPU models are placed on "cuda" during app startup. In a ZeroGPU Space,
#   Hugging Face's CUDA emulation handles this even though a real GPU is only
#   assigned while a @spaces.GPU function is executing.
#
# CPU path:
#   A separate float32 copy is loaded lazily only if somebody presses the CPU
#   button. This avoids paying the RAM cost unless CPU fallback is actually used.
# -----------------------------------------------------------------------------
_gpu_tokenizer = None
_gpu_encoder = None
_gpu_repeat_model = None
_gpu_propeller_model = None
_gpu_init_error = None

_cpu_lock = threading.Lock()
_cpu_tokenizer = None
_cpu_encoder = None
_cpu_repeat_model = None
_cpu_propeller_model = None
_cpu_init_error = None


def _ensure_tread_source():
    """Fetch exactly the same TREAD commit used by the public Colab demo."""
    package_init = TREAD_DIR / "tread" / "__init__.py"
    if package_init.is_file():
        if str(TREAD_DIR) not in sys.path:
            sys.path.insert(0, str(TREAD_DIR))
        return

    shutil.rmtree(TREAD_DIR, ignore_errors=True)
    subprocess.run(["git", "init", str(TREAD_DIR)], check=True)
    subprocess.run(
        ["git", "-C", str(TREAD_DIR), "remote", "add", "origin", TREAD_REPO],
        check=True,
    )
    subprocess.run(
        [
            "git",
            "-C",
            str(TREAD_DIR),
            "fetch",
            "--depth",
            "1",
            "origin",
            TREAD_COMMIT,
        ],
        check=True,
    )
    subprocess.run(
        ["git", "-C", str(TREAD_DIR), "checkout", "--detach", "FETCH_HEAD"],
        check=True,
    )

    if not package_init.is_file():
        raise FileNotFoundError(
            f"TREAD installation incomplete: {package_init} not found"
        )

    if str(TREAD_DIR) not in sys.path:
        sys.path.insert(0, str(TREAD_DIR))


def _build_repeat_model(DMDModel, device):
    model = DMDModel(
        per_resi_emb_dim=1024,
        out_channel=64,
        hidden_dim=64,
        num_block=2,
        dropout=0.2,
        bilstm=True,
        kernel_size_conv1=3,
        kernel_size_block=7,
        multi=True,
        num_types=6,
        device=device,
    )
    state = torch.load(
        TREAD_DIR / "trained_model" / "linear-edge_model_repeatsdb.pt",
        map_location="cpu",
        weights_only=True,
    )
    model.load_state_dict(state)
    model = model.to(device).eval()
    return model


def _build_propeller_model(DMDModel, device):
    model = DMDModel(
        per_resi_emb_dim=1024,
        out_channel=128,
        hidden_dim=64,
        num_block=1,
        dropout=0.2,
        bilstm=True,
        kernel_size_conv1=11,
        kernel_size_block=7,
        device=device,
    )
    state = torch.load(
        TREAD_DIR / "trained_model" / "linear-edge_model_propeller_blade.pt",
        map_location="cpu",
        weights_only=True,
    )
    model.load_state_dict(state)
    model = model.to(device).eval()
    return model


def initialize_gpu_models_at_startup():
    """Prepare the ZeroGPU model copy on emulated CUDA during app startup."""
    global _gpu_tokenizer, _gpu_encoder, _gpu_repeat_model
    global _gpu_propeller_model, _gpu_init_error

    try:
        _ensure_tread_source()
        from tread.model import DMDModel

        _gpu_tokenizer = T5Tokenizer.from_pretrained(
            PROTT5_MODEL,
            do_lower_case=False,
            legacy=True,
        )
        _gpu_encoder = T5EncoderModel.from_pretrained(PROTT5_MODEL)
        _gpu_encoder = _gpu_encoder.to("cuda").eval()

        _gpu_repeat_model = _build_repeat_model(DMDModel, torch.device("cuda"))
        _gpu_propeller_model = _build_propeller_model(
            DMDModel, torch.device("cuda")
        )

        print("ZeroGPU model copy initialized successfully.")

    except Exception as exc:
        # Do not kill the whole web app: CPU fallback may still work and the
        # error will be shown clearly if the ZeroGPU button is pressed.
        _gpu_init_error = (
            f"{type(exc).__name__}: {exc}"
        )
        print("ZeroGPU initialization failed:", _gpu_init_error)


def initialize_cpu_models():
    """Load a separate CPU copy only when the CPU fallback is requested."""
    global _cpu_tokenizer, _cpu_encoder, _cpu_repeat_model
    global _cpu_propeller_model, _cpu_init_error

    if all(
        x is not None
        for x in (
            _cpu_tokenizer,
            _cpu_encoder,
            _cpu_repeat_model,
            _cpu_propeller_model,
        )
    ):
        return

    if _cpu_init_error is not None:
        raise RuntimeError(_cpu_init_error)

    with _cpu_lock:
        if all(
            x is not None
            for x in (
                _cpu_tokenizer,
                _cpu_encoder,
                _cpu_repeat_model,
                _cpu_propeller_model,
            )
        ):
            return

        try:
            _ensure_tread_source()
            from tread.model import DMDModel

            # Tokenizers are device independent. Reuse the one already loaded
            # for ZeroGPU if available.
            _cpu_tokenizer = _gpu_tokenizer
            if _cpu_tokenizer is None:
                _cpu_tokenizer = T5Tokenizer.from_pretrained(
                    PROTT5_MODEL,
                    do_lower_case=False,
                    legacy=True,
                )

            # CPU inference uses float32 for broad PyTorch CPU compatibility.
            # low_cpu_mem_usage reduces peak RAM while loading this large model.
            _cpu_encoder = T5EncoderModel.from_pretrained(
                PROTT5_MODEL,
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True,
            )
            _cpu_encoder = _cpu_encoder.to("cpu").eval()

            _cpu_repeat_model = _build_repeat_model(
                DMDModel, torch.device("cpu")
            )
            _cpu_propeller_model = _build_propeller_model(
                DMDModel, torch.device("cpu")
            )

            print("CPU fallback model copy initialized successfully.")

        except Exception as exc:
            _cpu_init_error = f"{type(exc).__name__}: {exc}"
            raise RuntimeError(_cpu_init_error) from exc


def clean_sequence(raw_sequence: str, max_length: int) -> str:
    if raw_sequence is None:
        raise gr.Error("Please paste a protein sequence.")

    text = raw_sequence.strip()
    if not text:
        raise gr.Error("Please paste a protein sequence.")

    # Accept either a raw sequence or a single FASTA record.
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if lines and lines[0].startswith(">"):
        lines = [line for line in lines[1:] if not line.startswith(">")]

    sequence = "".join(lines)
    sequence = re.sub(r"\s+", "", sequence).upper()

    if not sequence:
        raise gr.Error("No amino-acid sequence was found.")

    # Standard 20 amino acids plus the symbols handled by the notebook's
    # ProtT5 preprocessing (U, Z, O, B -> X) and X itself.
    allowed = set("ACDEFGHIKLMNPQRSTVWYUZOBX")
    invalid = sorted(set(sequence) - allowed)
    if invalid:
        raise gr.Error(
            "Invalid sequence characters: "
            + ", ".join(invalid)
            + ". Paste amino-acid letters only (a FASTA header is allowed)."
        )

    if len(sequence) > max_length:
        raise gr.Error(
            f"This backend currently accepts up to {max_length} residues; "
            f"your sequence has {len(sequence)} residues."
        )

    return sequence


def get_prott5_embedding(sequence, tokenizer, encoder, device):
    processed = " ".join(list(re.sub(r"[UZOB]", "X", sequence)))
    ids = tokenizer(
        [processed],
        add_special_tokens=True,
        padding="longest",
        truncation=False,
        return_attention_mask=True,
    )

    input_ids = torch.tensor(ids["input_ids"], device=device)
    attention_mask = torch.tensor(ids["attention_mask"], device=device)

    with torch.inference_mode():
        embedding_repr = encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )

    # Drop the terminal special token, exactly as in the working Colab demo.
    emb = embedding_repr.last_hidden_state[0, :-1]

    if emb.shape[0] != len(sequence):
        raise RuntimeError(
            f"ProtT5 returned {emb.shape[0]} residue embeddings for a "
            f"{len(sequence)}-residue sequence."
        )

    return emb


def get_ranges(preds, cutoff1=0.5, min_len=15, cutoff2=0.5, frac2=0.5):
    """Same motif-range logic as tread.utils.get_ranges."""
    preds = np.asarray(preds).flatten()
    above_threshold = preds > cutoff1
    peaks = []

    for key, group in groupby(enumerate(above_threshold), key=lambda x: x[1]):
        if key:
            group = list(group)
            if len(group) >= min_len:
                beg = group[0][0]
                end = beg + len(group)
                if (
                    len(np.where(preds[beg:end] > cutoff2)[0]) / len(group)
                    >= frac2
                ):
                    peaks.append((beg, end))

    return peaks


def make_plot(scores, ranges, title, cutoff):
    scores = np.asarray(scores).flatten()
    x = np.arange(1, len(scores) + 1)

    fig, ax = plt.subplots(figsize=(10, 4.5), dpi=150)
    ax.plot(x, scores, linewidth=1.5)
    ax.axhline(
        cutoff,
        linestyle="--",
        linewidth=1.0,
        alpha=0.7,
        label=f"cutoff = {cutoff:g}",
    )
    for start0, end0 in ranges:
        ax.axvspan(start0 + 1, end0, alpha=0.18)

    ax.set_xlim(1, max(1, len(scores)))
    ax.set_ylim(0, 1)
    ax.set_xlabel("Residue")
    ax.set_ylabel("Residue score")
    ax.set_title(title)
    ax.legend(loc="upper right")
    fig.tight_layout()

    return fig


def make_download_table(sequence, scores, ranges):
    scores = np.asarray(scores).flatten()
    predicted = np.zeros(len(sequence), dtype=bool)

    for start0, end0 in ranges:
        predicted[start0:end0] = True

    return pd.DataFrame(
        {
            "residue": np.arange(1, len(sequence) + 1),
            "amino_acid": list(sequence),
            "score": scores,
            "predicted_motif": predicted,
        }
    )


def save_csv(df: pd.DataFrame):
    temp = NamedTemporaryFile(
        prefix="tread_prediction_",
        suffix=".csv",
        delete=False,
    )
    temp.close()
    df.to_csv(temp.name, index=False)
    return temp.name


def _predict_core(
    sequence_text,
    model_choice,
    repeat_head,
    cutoff,
    min_len,
    *,
    tokenizer,
    encoder,
    repeat_model,
    propeller_model,
    device,
    backend_label,
    max_length,
):
    sequence = clean_sequence(sequence_text, max_length=max_length)

    embedding = get_prott5_embedding(
        sequence,
        tokenizer=tokenizer,
        encoder=encoder,
        device=device,
    )

    with torch.inference_mode():
        if model_choice == "RepeatsDB repeat annotation":
            seg_prediction, type_prediction = repeat_model.predict_single(
                embedding
            )

            if repeat_head == "Any repeat (segmentation head)":
                scores = seg_prediction
            else:
                scores = type_prediction[TYPE_HEAD_INDEX[repeat_head]]

            profile_name = repeat_head
            title = f"TREAD RepeatsDB — {profile_name}"

        else:
            scores = propeller_model.predict_single(embedding)
            profile_name = "Beta-propeller blade"
            title = "TREAD — beta-propeller blade annotation"

    ranges = get_ranges(
        scores,
        cutoff1=float(cutoff),
        min_len=int(min_len),
    )
    fig = make_plot(scores, ranges, title, float(cutoff))

    # Convert notebook's 0-based, end-exclusive ranges to user-facing
    # 1-based inclusive residue coordinates.
    if ranges:
        ranges_df = pd.DataFrame(
            [
                {
                    "Start (1-based)": start0 + 1,
                    "End (1-based)": end0,
                    "Length": end0 - start0,
                }
                for start0, end0 in ranges
            ]
        )
        range_text = ", ".join(f"{s + 1}–{e}" for s, e in ranges)
    else:
        ranges_df = pd.DataFrame(
            columns=["Start (1-based)", "End (1-based)", "Length"]
        )
        range_text = "None at the current thresholds"

    profile_df = make_download_table(sequence, scores, ranges)
    csv_path = save_csv(profile_df)

    status = (
        f"**Sequence length:** {len(sequence)} aa  \n"
        f"**Model:** {model_choice}  \n"
        f"**Profile:** {profile_name}  \n"
        f"**Predicted motif ranges:** {range_text}  \n"
        f"**Compute backend:** {backend_label}"
    )

    return status, fig, ranges_df, csv_path


@spaces.GPU(duration=120)
def predict_gpu(sequence_text, model_choice, repeat_head, cutoff, min_len):
    """ZeroGPU path. A real GPU is allocated only for this function call."""
    if _gpu_init_error is not None:
        raise gr.Error(
            "The ZeroGPU model copy could not be initialized. "
            f"Startup error: {_gpu_init_error}"
        )

    if any(
        x is None
        for x in (
            _gpu_tokenizer,
            _gpu_encoder,
            _gpu_repeat_model,
            _gpu_propeller_model,
        )
    ):
        raise gr.Error(
            "ZeroGPU models are not initialized. Check the Space runtime log."
        )

    device_name = "ZeroGPU"
    try:
        device_name = f"ZeroGPU — {torch.cuda.get_device_name(0)}"
    except Exception:
        pass

    return _predict_core(
        sequence_text,
        model_choice,
        repeat_head,
        cutoff,
        min_len,
        tokenizer=_gpu_tokenizer,
        encoder=_gpu_encoder,
        repeat_model=_gpu_repeat_model,
        propeller_model=_gpu_propeller_model,
        device=torch.device("cuda"),
        backend_label=device_name,
        max_length=MAX_GPU_SEQUENCE_LENGTH,
    )


def predict_cpu(sequence_text, model_choice, repeat_head, cutoff, min_len):
    """CPU fallback. Does not consume ZeroGPU quota, but is much slower."""
    try:
        initialize_cpu_models()
    except Exception as exc:
        raise gr.Error(f"CPU model initialization failed: {exc}") from exc

    return _predict_core(
        sequence_text,
        model_choice,
        repeat_head,
        cutoff,
        min_len,
        tokenizer=_cpu_tokenizer,
        encoder=_cpu_encoder,
        repeat_model=_cpu_repeat_model,
        propeller_model=_cpu_propeller_model,
        device=torch.device("cpu"),
        backend_label="CPU fallback",
        max_length=MAX_CPU_SEQUENCE_LENGTH,
    )


def load_repeat_example():
    return (
        EXAMPLE_REPEAT,
        "RepeatsDB repeat annotation",
        "Beta-barrel",
        0.8,
        20,
    )


def load_propeller_example():
    return (
        EXAMPLE_PROPELLER,
        "Beta-propeller blade annotation",
        "Beta-propeller",
        0.8,
        20,
    )


# IMPORTANT for ZeroGPU:
# Prepare the CUDA/emulated-CUDA model copy at module level, before requests.
initialize_gpu_models_at_startup()


with gr.Blocks(title="TREAD — Protein Repeat Annotation") as demo:
    gr.Markdown(
        """
# TREAD — Protein Repeat Annotation

Paste a protein sequence and run one of the two pretrained TREAD models.

- **RepeatsDB repeat annotation:** residue-wise repeat segmentation plus six repeat-fold heads.
- **Beta-propeller blade annotation:** residue-wise blade prediction.

**Recommended:** use **Run with ZeroGPU**.  
**Fallback:** use **Run on CPU** if GPU quota/availability is a problem; CPU inference is substantially slower.

ProtT5 embeddings are generated on the fly with `Rostlab/prot_t5_xl_half_uniref50-enc`.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            sequence_input = gr.Textbox(
                label="Protein sequence",
                lines=10,
                placeholder=(
                    "Paste a raw amino-acid sequence or a single FASTA record..."
                ),
            )
            model_choice = gr.Radio(
                choices=[
                    "RepeatsDB repeat annotation",
                    "Beta-propeller blade annotation",
                ],
                value="RepeatsDB repeat annotation",
                label="Model",
            )
            repeat_head = gr.Dropdown(
                choices=REPEAT_HEADS,
                value="Beta-barrel",
                label="RepeatsDB profile to display",
                info=(
                    "Used only for the RepeatsDB model. "
                    "The default matches the public Colab example."
                ),
            )

        with gr.Column(scale=2):
            cutoff = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.8,
                step=0.01,
                label="Residue score threshold",
            )
            min_len = gr.Slider(
                minimum=1,
                maximum=100,
                value=20,
                step=1,
                label="Minimum motif length",
            )

            gpu_button = gr.Button(
                "Run with ZeroGPU (recommended)",
                variant="primary",
            )
            cpu_button = gr.Button(
                "Run on CPU (slow fallback)",
                variant="secondary",
            )

            with gr.Row():
                repeat_example_button = gr.Button("Load RepeatsDB example")
                propeller_example_button = gr.Button("Load propeller example")

    gr.Markdown("## Results")
    status_output = gr.Markdown()
    plot_output = gr.Plot(label="Residue-wise score profile")
    ranges_output = gr.Dataframe(
        headers=["Start (1-based)", "End (1-based)", "Length"],
        label="Predicted motif ranges",
        interactive=False,
    )
    download_output = gr.File(label="Download per-residue CSV")

    gr.Markdown(
        f"""
### Notes

- ZeroGPU accepts sequences up to **{MAX_GPU_SEQUENCE_LENGTH} aa** in this web demo.
- CPU fallback accepts sequences up to **{MAX_CPU_SEQUENCE_LENGTH} aa** by default because ProtT5-XL is very slow on the free CPU backend.
- The range-calling logic uses the same thresholding rule as the public Colab/TREAD utility.
- The web table reports **1-based inclusive** residue coordinates for readability.
- Rare/ambiguous residues `U`, `Z`, `O`, and `B` are mapped to `X` for ProtT5 embedding, matching the Colab demo.

Source code and pretrained TREAD checkpoints: [KYQiu21/TREAD](https://github.com/KYQiu21/TREAD)
        """
    )

    gpu_button.click(
        fn=predict_gpu,
        inputs=[
            sequence_input,
            model_choice,
            repeat_head,
            cutoff,
            min_len,
        ],
        outputs=[
            status_output,
            plot_output,
            ranges_output,
            download_output,
        ],
    )

    cpu_button.click(
        fn=predict_cpu,
        inputs=[
            sequence_input,
            model_choice,
            repeat_head,
            cutoff,
            min_len,
        ],
        outputs=[
            status_output,
            plot_output,
            ranges_output,
            download_output,
        ],
    )

    repeat_example_button.click(
        fn=load_repeat_example,
        inputs=[],
        outputs=[
            sequence_input,
            model_choice,
            repeat_head,
            cutoff,
            min_len,
        ],
    )

    propeller_example_button.click(
        fn=load_propeller_example,
        inputs=[],
        outputs=[
            sequence_input,
            model_choice,
            repeat_head,
            cutoff,
            min_len,
        ],
    )


demo.queue()

if __name__ == "__main__":
    demo.launch()