Spaces:

hugging-apps
/

nvidia-rfd4-proteina-diffusers

Running on Zero

Upload folder using huggingface_hub

f9d3b3a verified 4 days ago

11.7 kB

import os

# Environment preamble before any heavy import: these gate AtomWorks/numba/biotite startup and
# keep the ZeroGPU parent process free of CUDA initialization.
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
os.environ.setdefault("ALLOW_BIOTITE_CCD", "true")
os.environ.setdefault("DISABLE_CUEQUIVARIANCE", "true")
os.environ.setdefault("WANDB_MODE", "disabled")
os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")
os.environ.setdefault("CLUSTER", "local")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

import spaces  # noqa: E402  # MUST precede torch / diffusers / atomworks

import time
from pathlib import Path

import gradio as gr
import torch

MODEL_ID = "nvidia/RFD4-Proteina-Diffusers"
EXAMPLES_DIR = Path(__file__).parent / "examples"

from diffusers import RFD4ProteinaPipeline  # noqa: E402
from rfd4_proteina_prep import prepare_cif, write_cif  # noqa: E402

# The feature tensors the pipeline consumes as direct keyword arguments; everything else the
# preprocessor emits rides inside `conditioning_features`.
REQUIRED_KEYS = (
    "padding_mask",
    "is_in_context_conditioning",
    "supports_local_latent",
    "token_type",
    "atom37_mask",
)

MODES = (
    "Full design",
    "Inverse folding",
    "Partial diffusion",
)

MODE_INFO = {
    "Full design": (
        "Design every free token in the input: sequence and structure are generated from noise, "
        "conditioned on whatever the annotated CIF holds fixed."
    ),
    "Inverse folding": (
        "The backbone of every free token is held fixed; RFD4 predicts only the sequence. Use it on "
        "a structure whose fold you want to keep."
    ),
    "Partial diffusion": (
        "Start from the input structure at a partly-noised state: the beginning of the trajectory "
        "is skipped, so designs stay close to the input fold."
    ),
}

# The released recipe fixes the number of diffusion steps (len(config.recycle_depths) = 60).
N_STEPS = 60

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

_PLACEHOLDER_HTML = (
    '<div style="width:100%;height:520px;display:flex;align-items:center;justify-content:center;'
    "color:var(--body-text-color);border:1px dashed var(--border-color-primary);border-radius:8px;"
    'font-size:14px;">Press “Design” to see the generated structure here</div>'
)

# The pipeline is loaded once at module scope, eagerly on the GPU: no offloading, no device_map.
pipe = RFD4ProteinaPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")


def _describe(mode: str) -> str:
    """One-line explanation of a design mode, shown under the dropdown."""
    return MODE_INFO.get(mode, "")


def _viewer_for(cif_path: str | Path, height: int = 520) -> str:
    """Render a designed CIF as an interactive 3Dmol.js cartoon view."""
    import py3Dmol

    with open(cif_path) as handle:
        data = handle.read()
    view = py3Dmol.view(width="100%", height=height)
    view.addModel(data, "cif")
    view.setStyle({"cartoon": {"colorscheme": "chain"}})
    view.zoomTo()
    return view._make_html()


@spaces.GPU(duration=120)
def design_protein(
    cif_file,
    mode="Full design",
    seed=42,
    num_designs=1,
    remaining_fraction=0.4,
    progress=gr.Progress(track_tqdm=True),
):
    """Run RFD4-Proteina on an annotated CIF and return the designed structure(s).
    Args:
        cif_file: Path to the uploaded annotated CIF (AtomWorks Studio format).
        mode: "Full design", "Inverse folding" or "Partial diffusion".
        seed: Random seed; the same input and seed reproduce the same design.
        num_designs: How many independent designs to draw in one call (1-8).
        remaining_fraction: "Partial diffusion" only — the fraction of the trajectory left to
            run. The released example uses 0.4.
        progress: Gradio progress bar.
    Returns:
        The design CIF path (or list of paths), an interactive 3D view of the first design, and a
        short text summary.
    """
    started = time.perf_counter()
    if cif_file is None:
        raise gr.Error("Upload an annotated CIF first, or pick one of the examples below.")
    num_designs = int(num_designs)
    if not 1 <= num_designs <= 8:
        raise gr.Error("Number of designs must be between 1 and 8.")

    # Feature preparation is CPU-side AtomWorks parsing; the pipeline moves every tensor to its
    # execution device itself, so CPU tensors are fine here.
    features, atoms = prepare_cif(str(cif_file), seed=int(seed), device="cpu")
    tokens = int(features["padding_mask"].shape[1])

    required = {name: features[name] for name in REQUIRED_KEYS}
    conditioning = {name: value for name, value in features.items() if name not in REQUIRED_KEYS}

    if mode == "Inverse folding":
        kwargs = {
            "inverse_folding": True,
            "partial_t_bb_ca": 0.0,
            "partial_t_local_latents": 0.0,
        }
        required["coords_nm"] = features["coords_nm"]
    elif mode == "Partial diffusion":
        # Reproduce the released partial-diffusion entry point: start each modality at the time
        # its schedule reaches after skipping (1 - remaining_fraction) of the trajectory.
        fraction = float(min(max(remaining_fraction, 0.05), 1.0))
        index = round((1 - fraction) * N_STEPS)
        pipe.bb_ca_scheduler.set_timesteps(N_STEPS)
        pipe.local_latents_scheduler.set_timesteps(N_STEPS)
        kwargs = {
            "partial_t_bb_ca": float(pipe.bb_ca_scheduler.timesteps[index]),
            "partial_t_local_latents": float(pipe.local_latents_scheduler.timesteps[index]),
        }
        required["coords_nm"] = features["coords_nm"]
    else:
        kwargs = {}
        conditioning["coords_nm"] = features["coords_nm"]

    if num_designs == 1:
        generator = torch.Generator("cuda").manual_seed(int(seed))
    else:
        generator = [
            torch.Generator("cuda").manual_seed(int(seed) + offset)
            for offset in range(num_designs)
        ]

    with torch.no_grad():
        output = pipe(
            padding_mask=required["padding_mask"],
            is_in_context_conditioning=required["is_in_context_conditioning"],
            supports_local_latent=required["supports_local_latent"],
            token_type=required["token_type"],
            atom37_mask=required["atom37_mask"],
            coords_nm=required.get("coords_nm"),
            conditioning_features=conditioning,
            num_designs_per_prompt=num_designs,
            generator=generator,
            **kwargs,
        )

    # `write_cif` reads one feature row per output row: repeat the input features per design.
    if num_designs > 1:
        features = {
            name: value.repeat_interleave(num_designs, dim=0) for name, value in features.items()
        }

    out_paths = []
    for row in range(num_designs):
        out_paths.append(write_cif(output, features, atoms, _unique_path(row, seed), row=row))

    elapsed = time.perf_counter() - started
    summary = (
        f"**{num_designs}** design{'s' if num_designs > 1 else ''} · **{tokens}** tokens · "
        f"**{mode.lower()}** · seed **{int(seed)}** · **{elapsed:.1f} s** on this GPU"
    )
    viewer_html = _viewer_for(out_paths[0])

    if num_designs == 1:
        return out_paths[0], viewer_html, summary
    return out_paths, viewer_html, summary


def _unique_path(row: int, seed) -> Path:
    """A fresh, non-existent output path in this worker's temp directory."""
    base = Path(os.environ.get("GRADIO_TEMP_DIR") or "/tmp/gradio") / f"rfd4_design_{int(seed)}_{row}.cif"
    base.parent.mkdir(parents=True, exist_ok=True)
    candidate = base
    suffix = 0
    while candidate.exists():
        suffix += 1
        candidate = base.with_name(f"{base.stem}_{suffix}{base.suffix}")
    return candidate


with gr.Blocks(title="RFD4-Proteina · Diffusers") as demo:
    gr.HTML(
        """
    <h1 style="margin: 0; font-size: 28px;">RFD4-Proteina · Diffusers</h1>
    <p style="margin-top: 4px;">
      Joint sequence–structure generation for biomolecular design, conditioned on an
      <b>annotated CIF</b> that says what to hold fixed and what to design. This Space runs
      <code>RFD4ProteinaPipeline</code> from 🤗 Diffusers on ZeroGPU.
    </p>
    """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=1):
                cif_input = gr.File(
                    label="Annotated CIF (input)",
                    file_types=[".cif"],
                    type="filepath",
                    file_count="single",
                )
                mode = gr.Dropdown(
                    choices=list(MODES),
                    value="Full design",
                    label="Design mode",
                    info="What RFD4 generates versus keeps from the input.",
                )
                mode_info = gr.Markdown(_describe("Full design"))
                with gr.Accordion("Advanced options", open=False):
                    seed = gr.Slider(0, 2**31 - 1, value=42, step=1, label="Seed")
                    num_designs = gr.Slider(1, 8, value=1, step=1, label="Number of designs")
                    remaining_fraction = gr.Slider(
                        0.05,
                        1.0,
                        value=0.4,
                        step=0.05,
                        label="Remaining schedule (partial diffusion only)",
                        info="The released partial-diffusion example uses 0.4.",
                    )
                go = gr.Button("Design", variant="primary")
            with gr.Column(scale=2):
                viewer = gr.HTML(value=_PLACEHOLDER_HTML, show_label=False)
                out_files = gr.File(
                    label="Designed structure(s) · annotated CIF",
                    file_count="single",
                    interactive=False,
                )
                summary = gr.Markdown()
        mode.change(_describe, inputs=mode, outputs=mode_info)
        go.click(
            design_protein,
            inputs=[cif_input, mode, seed, num_designs, remaining_fraction],
            outputs=[out_files, viewer, summary],
            api_name="design",
        )

        gr.Examples(
            examples=[
                [str(EXAMPLES_DIR / "motif_scaffolding.cif"), "Full design"],
                [str(EXAMPLES_DIR / "inverse_folding.cif"), "Inverse folding"],
                [str(EXAMPLES_DIR / "inverse_folding.cif"), "Partial diffusion"],
                [str(EXAMPLES_DIR / "protein_binder.cif"), "Full design"],
                [str(EXAMPLES_DIR / "enzyme_design.cif"), "Full design"],
                [str(EXAMPLES_DIR / "small_molecule_binder.cif"), "Full design"],
                [str(EXAMPLES_DIR / "unconditional.cif"), "Full design"],
            ],
            inputs=[cif_input, mode],
            outputs=[out_files, viewer, summary],
            fn=design_protein,
            cache_examples=True,
            cache_mode="lazy",
            example_labels=[
                "Motif scaffolding (1QJG)",
                "Inverse folding (2KL8)",
                "Partial diffusion (2KL8)",
                "Protein binder (PD-L1)",
                "Enzyme (catalytic motif)",
                "Small-molecule binder (SAM)",
                "Unconditional protein",
            ],
            label="Author examples (from the model card, BSD-3-Clause)",
        )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)

	import os

	# Environment preamble before any heavy import: these gate AtomWorks/numba/biotite startup and
	# keep the ZeroGPU parent process free of CUDA initialization.
	os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
	os.environ.setdefault("ALLOW_BIOTITE_CCD", "true")
	os.environ.setdefault("DISABLE_CUEQUIVARIANCE", "true")
	os.environ.setdefault("WANDB_MODE", "disabled")
	os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false")
	os.environ.setdefault("CLUSTER", "local")
	os.environ.setdefault("OMP_NUM_THREADS", "1")
	os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

	import spaces # noqa: E402 # MUST precede torch / diffusers / atomworks

	import time
	from pathlib import Path

	import gradio as gr
	import torch

	MODEL_ID = "nvidia/RFD4-Proteina-Diffusers"
	EXAMPLES_DIR = Path(__file__).parent / "examples"

	from diffusers import RFD4ProteinaPipeline # noqa: E402
	from rfd4_proteina_prep import prepare_cif, write_cif # noqa: E402

	# The feature tensors the pipeline consumes as direct keyword arguments; everything else the
	# preprocessor emits rides inside `conditioning_features`.
	REQUIRED_KEYS = (
	"padding_mask",
	"is_in_context_conditioning",
	"supports_local_latent",
	"token_type",
	"atom37_mask",
	)

	MODES = (
	"Full design",
	"Inverse folding",
	"Partial diffusion",
	)

	MODE_INFO = {
	"Full design": (
	"Design every free token in the input: sequence and structure are generated from noise, "
	"conditioned on whatever the annotated CIF holds fixed."
	),
	"Inverse folding": (
	"The backbone of every free token is held fixed; RFD4 predicts only the sequence. Use it on "
	"a structure whose fold you want to keep."
	),
	"Partial diffusion": (
	"Start from the input structure at a partly-noised state: the beginning of the trajectory "
	"is skipped, so designs stay close to the input fold."
	),
	}

	# The released recipe fixes the number of diffusion steps (len(config.recycle_depths) = 60).
	N_STEPS = 60

	CSS = """
	#col-container { max-width: 1100px; margin: 0 auto; }
	.dark .gradio-container { color: var(--body-text-color); }
	"""

	_PLACEHOLDER_HTML = (
	'<div style="width:100%;height:520px;display:flex;align-items:center;justify-content:center;'
	"color:var(--body-text-color);border:1px dashed var(--border-color-primary);border-radius:8px;"
	'font-size:14px;">Press “Design” to see the generated structure here</div>'
	)

	# The pipeline is loaded once at module scope, eagerly on the GPU: no offloading, no device_map.
	pipe = RFD4ProteinaPipeline.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")


	def _describe(mode: str) -> str:
	"""One-line explanation of a design mode, shown under the dropdown."""
	return MODE_INFO.get(mode, "")


	def _viewer_for(cif_path: str \| Path, height: int = 520) -> str:
	"""Render a designed CIF as an interactive 3Dmol.js cartoon view."""
	import py3Dmol

	with open(cif_path) as handle:
	data = handle.read()
	view = py3Dmol.view(width="100%", height=height)
	view.addModel(data, "cif")
	view.setStyle({"cartoon": {"colorscheme": "chain"}})
	view.zoomTo()
	return view._make_html()


	@spaces.GPU(duration=120)
	def design_protein(
	cif_file,
	mode="Full design",
	seed=42,
	num_designs=1,
	remaining_fraction=0.4,
	progress=gr.Progress(track_tqdm=True),
	):
	"""Run RFD4-Proteina on an annotated CIF and return the designed structure(s).

	Args:
	cif_file: Path to the uploaded annotated CIF (AtomWorks Studio format).
	mode: "Full design", "Inverse folding" or "Partial diffusion".
	seed: Random seed; the same input and seed reproduce the same design.
	num_designs: How many independent designs to draw in one call (1-8).
	remaining_fraction: "Partial diffusion" only — the fraction of the trajectory left to
	run. The released example uses 0.4.
	progress: Gradio progress bar.

	Returns:
	The design CIF path (or list of paths), an interactive 3D view of the first design, and a
	short text summary.
	"""
	started = time.perf_counter()
	if cif_file is None:
	raise gr.Error("Upload an annotated CIF first, or pick one of the examples below.")
	num_designs = int(num_designs)
	if not 1 <= num_designs <= 8:
	raise gr.Error("Number of designs must be between 1 and 8.")

	# Feature preparation is CPU-side AtomWorks parsing; the pipeline moves every tensor to its
	# execution device itself, so CPU tensors are fine here.
	features, atoms = prepare_cif(str(cif_file), seed=int(seed), device="cpu")
	tokens = int(features["padding_mask"].shape[1])

	required = {name: features[name] for name in REQUIRED_KEYS}
	conditioning = {name: value for name, value in features.items() if name not in REQUIRED_KEYS}

	if mode == "Inverse folding":
	kwargs = {
	"inverse_folding": True,
	"partial_t_bb_ca": 0.0,
	"partial_t_local_latents": 0.0,
	}
	required["coords_nm"] = features["coords_nm"]
	elif mode == "Partial diffusion":
	# Reproduce the released partial-diffusion entry point: start each modality at the time
	# its schedule reaches after skipping (1 - remaining_fraction) of the trajectory.
	fraction = float(min(max(remaining_fraction, 0.05), 1.0))
	index = round((1 - fraction) * N_STEPS)
	pipe.bb_ca_scheduler.set_timesteps(N_STEPS)
	pipe.local_latents_scheduler.set_timesteps(N_STEPS)
	kwargs = {
	"partial_t_bb_ca": float(pipe.bb_ca_scheduler.timesteps[index]),
	"partial_t_local_latents": float(pipe.local_latents_scheduler.timesteps[index]),
	}
	required["coords_nm"] = features["coords_nm"]
	else:
	kwargs = {}
	conditioning["coords_nm"] = features["coords_nm"]

	if num_designs == 1:
	generator = torch.Generator("cuda").manual_seed(int(seed))
	else:
	generator = [
	torch.Generator("cuda").manual_seed(int(seed) + offset)
	for offset in range(num_designs)
	]

	with torch.no_grad():
	output = pipe(
	padding_mask=required["padding_mask"],
	is_in_context_conditioning=required["is_in_context_conditioning"],
	supports_local_latent=required["supports_local_latent"],
	token_type=required["token_type"],
	atom37_mask=required["atom37_mask"],
	coords_nm=required.get("coords_nm"),
	conditioning_features=conditioning,
	num_designs_per_prompt=num_designs,
	generator=generator,
	**kwargs,
	)

	# `write_cif` reads one feature row per output row: repeat the input features per design.
	if num_designs > 1:
	features = {
	name: value.repeat_interleave(num_designs, dim=0) for name, value in features.items()
	}

	out_paths = []
	for row in range(num_designs):
	out_paths.append(write_cif(output, features, atoms, _unique_path(row, seed), row=row))

	elapsed = time.perf_counter() - started
	summary = (
	f"{num_designs} design{'s' if num_designs > 1 else ''} · {tokens} tokens · "
	f"{mode.lower()} · seed {int(seed)} · {elapsed:.1f} s on this GPU"
	)
	viewer_html = _viewer_for(out_paths[0])

	if num_designs == 1:
	return out_paths[0], viewer_html, summary
	return out_paths, viewer_html, summary


	def _unique_path(row: int, seed) -> Path:
	"""A fresh, non-existent output path in this worker's temp directory."""
	base = Path(os.environ.get("GRADIO_TEMP_DIR") or "/tmp/gradio") / f"rfd4_design_{int(seed)}_{row}.cif"
	base.parent.mkdir(parents=True, exist_ok=True)
	candidate = base
	suffix = 0
	while candidate.exists():
	suffix += 1
	candidate = base.with_name(f"{base.stem}_{suffix}{base.suffix}")
	return candidate


	with gr.Blocks(title="RFD4-Proteina · Diffusers") as demo:
	gr.HTML(
	"""
	<h1 style="margin: 0; font-size: 28px;">RFD4-Proteina · Diffusers</h1>
	<p style="margin-top: 4px;">
	Joint sequence–structure generation for biomolecular design, conditioned on an
	<b>annotated CIF</b> that says what to hold fixed and what to design. This Space runs
	<code>RFD4ProteinaPipeline</code> from 🤗 Diffusers on ZeroGPU.
	</p>
	"""
	)
	with gr.Column(elem_id="col-container"):
	with gr.Row():
	with gr.Column(scale=1):
	cif_input = gr.File(
	label="Annotated CIF (input)",
	file_types=[".cif"],
	type="filepath",
	file_count="single",
	)
	mode = gr.Dropdown(
	choices=list(MODES),
	value="Full design",
	label="Design mode",
	info="What RFD4 generates versus keeps from the input.",
	)
	mode_info = gr.Markdown(_describe("Full design"))
	with gr.Accordion("Advanced options", open=False):
	seed = gr.Slider(0, 2**31 - 1, value=42, step=1, label="Seed")
	num_designs = gr.Slider(1, 8, value=1, step=1, label="Number of designs")
	remaining_fraction = gr.Slider(
	0.05,
	1.0,
	value=0.4,
	step=0.05,
	label="Remaining schedule (partial diffusion only)",
	info="The released partial-diffusion example uses 0.4.",
	)
	go = gr.Button("Design", variant="primary")
	with gr.Column(scale=2):
	viewer = gr.HTML(value=_PLACEHOLDER_HTML, show_label=False)
	out_files = gr.File(
	label="Designed structure(s) · annotated CIF",
	file_count="single",
	interactive=False,
	)
	summary = gr.Markdown()
	mode.change(_describe, inputs=mode, outputs=mode_info)
	go.click(
	design_protein,
	inputs=[cif_input, mode, seed, num_designs, remaining_fraction],
	outputs=[out_files, viewer, summary],
	api_name="design",
	)

	gr.Examples(
	examples=[
	[str(EXAMPLES_DIR / "motif_scaffolding.cif"), "Full design"],
	[str(EXAMPLES_DIR / "inverse_folding.cif"), "Inverse folding"],
	[str(EXAMPLES_DIR / "inverse_folding.cif"), "Partial diffusion"],
	[str(EXAMPLES_DIR / "protein_binder.cif"), "Full design"],
	[str(EXAMPLES_DIR / "enzyme_design.cif"), "Full design"],
	[str(EXAMPLES_DIR / "small_molecule_binder.cif"), "Full design"],
	[str(EXAMPLES_DIR / "unconditional.cif"), "Full design"],
	],
	inputs=[cif_input, mode],
	outputs=[out_files, viewer, summary],
	fn=design_protein,
	cache_examples=True,
	cache_mode="lazy",
	example_labels=[
	"Motif scaffolding (1QJG)",
	"Inverse folding (2KL8)",
	"Partial diffusion (2KL8)",
	"Protein binder (PD-L1)",
	"Enzyme (catalytic motif)",
	"Small-molecule binder (SAM)",
	"Unconditional protein",
	],
	label="Author examples (from the model card, BSD-3-Clause)",
	)

	demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)