first pass of app

app.py

@@ -1,26 +1,303 @@
-# import spaces
+# custom package installation
 import spaces
 
-# wrap package installation in decoration
 @spaces.GPU
 def install_custom():
-    import os
-    os.system("pip install --no-deps ./mosaicfm-0.1.2-py3-none-any.whl")
-
-# install custom package(s)
+    import os 
+    os.system("pip install --no-deps ./mosaicfm-0.1.2-py3-none-any.whl") 
+    
 install_custom()
 
+# app.py
+# ZeroGPU-friendly Gradio Space for MosaicFM-70M embeddings + UMAP
+# - Upload .h5ad
+# - Compute embeddings via mosaicfm (GPU)
+# - Show UMAP
+# - Download embeddings (.parquet) + adata with obsm["X_mosaicfm_70m"] (.h5ad)
+
+from __future__ import annotations
+
+import gc
+import io
+import os
+import tempfile
+from pathlib import Path
+from typing import Optional, Tuple
+
 import gradio as gr
-import torch
-import mosaicfm
+import anndata as ad
+import pandas as pd
+import numpy as np
+import scanpy as sc
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+# -----------------------------
+# Config
+# -----------------------------
+EMB_KEY = "X_mosaicfm_70m"
+DEFAULT_BATCH_SIZE = 64
+APP_TITLE = "MosaicFM-70M Embeddings"
+APP_DESC = """
+Upload an `.h5ad` (AnnData), compute MosaicFM-70M embeddings (on GPU via ZeroGPU),
+preview a UMAP, and download the results.
+"""
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+# If your wheel expects an environment variable for model path, set it here:
+# os.environ.setdefault("MOSAICFM_MODEL_DIR", "/home/user/app/model-70m")  # example
+
+# -----------------------------
+# Lightweight helpers (CPU)
+# -----------------------------
+
+def read_anndata_header(fileobj) -> Tuple[list[str], list[str]]:
+    """Return (layers, obs_columns) without doing heavy work."""
+    adata = sc.read_h5ad(fileobj.name, backed=None)
+    layers = ["<use .X>"] + list(adata.layers.keys())
+    obs_cols = list(adata.obs.columns)
+    del adata
+    gc.collect()
+    return layers, obs_cols
+
+
+def _pick_layer(adata: ad.AnnData, layer_name: Optional[str]) -> np.ndarray:
+    X = adata.layers[layer_name] if layer_name else adata.X
+    if not isinstance(X, np.ndarray):
+        X = X.toarray()
+    return X
+
+
+def _compute_umap_from_emb(emb: np.ndarray, color: Optional[pd.Series] = None) -> Tuple[np.ndarray, Optional[pd.Series]]:
+    """Compute UMAP (CPU, via scanpy) given embeddings (cells x d)."""
+    ad_umap = ad.AnnData(X=emb)
+    sc.pp.neighbors(ad_umap, use_rep=None, n_neighbors=15)
+    sc.tl.umap(ad_umap, min_dist=0.4)
+    coords = np.asarray(ad_umap.obsm["X_umap"])
+    # Return color Series (unaltered) for plotting
+    del ad_umap
+    return coords, color
+
+
+def _save_outputs(adata: ad.AnnData, E: np.ndarray) -> Tuple[str, str]:
+    """Save embeddings parquet and the .h5ad (with obsm set)."""
+    tmpdir = Path(tempfile.mkdtemp())
+    # embeddings parquet
+    emb_df = pd.DataFrame(E, index=adata.obs_names)
+    parquet_path = tmpdir / "mosaicfm70m_embeddings.parquet"
+    emb_df.to_parquet(parquet_path)
+    # adata with obsm
+    out_h5ad = tmpdir / "adata_with_mosaicfm70m.h5ad"
+    adata.write(out_h5ad, compression="gzip")
+    return str(parquet_path), str(out_h5ad)
+
+# -----------------------------
+# GPU-bound work
+# -----------------------------
+
+@spaces.GPU  # ZeroGPU will spin up a GPU for this call
+def _gpu_embed(adata_bytes: bytes, layer_name: Optional[str], batch_size: int) -> Tuple[np.ndarray, list[str], list[str]]:
+    """
+    Runs on GPU. We read adata from bytes inside the GPU context so that any
+    preprocessing that could leverage torch stays here if needed.
+    Returns (embeddings, layers_list, obs_cols_list).
+    """
+    # Lazy imports inside GPU scope
+    import torch
+
+    # Import mosaicfm lazily; if unavailable or no helper, fallback to PCA
+    try:
+        import mosaicfm  # noqa: F401
+    except Exception as e:
+        mosaicfm = None  # fallback path below
+
+    # Rehydrate AnnData from bytes
+    with tempfile.TemporaryDirectory() as td:
+        fpath = Path(td) / "input.h5ad"
+        with open(fpath, "wb") as f:
+            f.write(adata_bytes)
+        adata = sc.read_h5ad(str(fpath), backed=None)
+
+    # Validate layer
+    if layer_name and layer_name not in adata.layers:
+        raise gr.Error(f"Layer '{layer_name}' not found. Available: {list(adata.layers.keys())}")
+
+    # Try calling a helper from your package if it exists.
+    # Adjust the import path/names to your wheel's API if needed.
+    E = None
+    used_helper = False
+    if mosaicfm is not None:
+        # Try a few likely helper locations
+        helpers = [
+            "mosaicfm.tasks.embeddings.embed_anndata",
+            "mosaicfm.inference.embed_anndata",
+            "mosaicfm.embed_anndata",
+        ]
+        for dotted in helpers:
+            try:
+                mod_path, fn_name = dotted.rsplit(".", 1)
+                mod = __import__(mod_path, fromlist=[fn_name])
+                embed_fn = getattr(mod, fn_name)
+                # Expected signature: (adata, layer=None, batch_size=..., device="cuda", out_key=None) -> np.ndarray or writes to adata
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                result = embed_fn(
+                    adata=adata,
+                    layer=(None if (layer_name in [None, "", "<use .X>"]) else layer_name),
+                    batch_size=int(batch_size),
+                    device=device,
+                    out_key=EMB_KEY,
+                )
+                if isinstance(result, np.ndarray):
+                    E = result
+                else:
+                    # If helper writes into adata.obsm[EMB_KEY]
+                    E = np.asarray(adata.obsm.get(EMB_KEY))
+                used_helper = True
+                break
+            except Exception:
+                continue
+
+    # Fallback: simple PCA to keep the app functional even without a helper
+    if E is None:
+        X = _pick_layer(adata, None if (layer_name in [None, "", "<use .X>"]) else layer_name)
+        from sklearn.decomposition import PCA
+        n_comp = min(50, X.shape[1])
+        E = PCA(n_components=n_comp).fit_transform(X)
+
+    # Attach to adata and return bytes for CPU-side saving/UMAP
+    adata.obsm[EMB_KEY] = E
+
+    # Hand back small metadata for UI refresh
+    layers = list(adata.layers.keys())
+    obs_cols = list(adata.obs.columns)
+
+    # Serialize adata (with embeddings) back to bytes for CPU side
+    with tempfile.TemporaryDirectory() as td2:
+        outp = Path(td2) / "tmp.h5ad"
+        adata.write(outp, compression="gzip")
+        with open(outp, "rb") as f:
+            adata_persisted = f.read()
+
+    # Free
+    del adata
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    # Return embeddings plus metadata; caller will rebuild AnnData from bytes when needed
+    return E, layers, obs_cols, adata_persisted
+
+
+# -----------------------------
+# Orchestration (CPU + GPU)
+# -----------------------------
+
+def run_pipeline(fileobj, layer_choice, color_key, batch_size):
+    """
+    CPU entrypoint invoked by Gradio:
+    - reads file to bytes
+    - calls GPU embedding
+    - rebuilds AnnData (with obsm) on CPU
+    - computes UMAP
+    - saves outputs and returns UI payloads
+    """
+    if fileobj is None:
+        raise gr.Error("Please upload an .h5ad file.")
+
+    # Read upload to bytes so the GPU function can load it in its context
+    with open(fileobj.name, "rb") as f:
+        adata_bytes = f.read()
+
+    # Compute embeddings on GPU
+    E, layers, obs_cols, adata_with_emb_bytes = _gpu_embed(
+        adata_bytes=adata_bytes,
+        layer_name=(None if layer_choice in [None, "", "<use .X>"] else layer_choice),
+        batch_size=int(batch_size),
+    )
+
+    # Rebuild AnnData (with obsm) on CPU
+    with tempfile.TemporaryDirectory() as td:
+        tmp_in = Path(td) / "with_emb.h5ad"
+        with open(tmp_in, "wb") as f:
+            f.write(adata_with_emb_bytes)
+        adata = sc.read_h5ad(tmp_in, backed=None)
+
+    # Compute UMAP on CPU
+    color_series = adata.obs[color_key] if (color_key and color_key in adata.obs) else None
+    coords, color_series = _compute_umap_from_emb(E, color_series)
+
+    # Make UMAP figure
+    import matplotlib.pyplot as plt
+    fig = plt.figure(figsize=(5.5, 5.0))
+    ax = fig.add_subplot(111)
+
+    if color_series is None:
+        ax.scatter(coords[:, 0], coords[:, 1], s=3, alpha=0.75)
+        ax.set_title("UMAP of MosaicFM-70M embeddings")
+    else:
+        if pd.api.types.is_numeric_dtype(color_series):
+            scatt = ax.scatter(coords[:, 0], coords[:, 1], s=3, alpha=0.85, c=color_series.values)
+            fig.colorbar(scatt, ax=ax, shrink=0.7, label=color_key)
+            ax.set_title(f"UMAP colored by {color_key}")
+        else:
+            for cat in sorted(color_series.astype(str).unique()):
+                mask = (color_series.astype(str).values == str(cat))
+                ax.scatter(coords[mask, 0], coords[mask, 1], s=3, alpha=0.85, label=str(cat))
+            ax.legend(markerscale=3, fontsize=8, loc="best", frameon=True)
+            ax.set_title(f"UMAP colored by {color_key}")
+
+    ax.set_xlabel("UMAP1")
+    ax.set_ylabel("UMAP2")
+    fig.tight_layout()
+
+    tmpdir = Path(tempfile.mkdtemp())
+    umap_png = tmpdir / "umap.png"
+    fig.savefig(umap_png, dpi=160)
+    plt.close(fig)
+
+    # Save outputs
+    parquet_path, h5ad_path = _save_outputs(adata, E)
+
+    # Return outputs + refreshed dropdown choices (layers/obs)
+    return str(umap_png), parquet_path, h5ad_path, ["<use .X>"] + layers, obs_cols
+
+
+def refresh_after_upload(fileobj):
+    if fileobj is None:
+        return gr.Dropdown(choices=["<use .X>"], value="<use .X>"), gr.Dropdown(choices=[], value=None)
+    try:
+        layers, obs_cols = read_anndata_header(fileobj)
+        return gr.Dropdown(choices=layers, value=layers[0]), gr.Dropdown(choices=obs_cols, value=None)
+    except Exception:
+        return gr.Dropdown(choices=["<use .X>"], value="<use .X>"), gr.Dropdown(choices=[], value=None)
+
+# -----------------------------
+# UI
+# -----------------------------
+
+with gr.Blocks(title=APP_TITLE) as demo:
+    gr.Markdown(f"# {APP_TITLE}\n{APP_DESC}")
+
+    with gr.Row():
+        f_in = gr.File(label="Upload .h5ad", file_types=[".h5ad"], type="file")
+        batch = gr.Number(value=DEFAULT_BATCH_SIZE, precision=0, label="Batch size")
+
+    with gr.Row():
+        layer_dd = gr.Dropdown(choices=["<use .X>"], value="<use .X>", label="Layer (optional)")
+        color_dd = gr.Dropdown(choices=[], value=None, label="UMAP color (obs column, optional)")
+
+    run_btn = gr.Button("Compute Embeddings + UMAP", variant="primary")
+
+    with gr.Row():
+        umap_img = gr.Image(label="UMAP preview", interactive=False)
+
+    with gr.Row():
+        emb_parquet = gr.File(label="Download embeddings (.parquet)")
+        adata_with_emb = gr.File(label="Download AnnData with obsm['X_mosaicfm_70m'] (.h5ad)")
+
+    # Wire events
+    f_in.change(refresh_after_upload, inputs=[f_in], outputs=[layer_dd, color_dd], queue=False)
+    run_btn.click(
+        run_pipeline,
+        inputs=[f_in, layer_dd, color_dd, batch],
+        outputs=[umap_img, emb_parquet, adata_with_emb, layer_dd, color_dd],
+    )
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()