Allow iframe embedding for HF Spaces

app.py

@@ -1,162 +1,3 @@
-"""
-CASCADE – Flask app for HF Spaces
-Mounted at /cascade_v1/
-"""
-
-import os
-import json
-import uuid
-
-import redis
-from flask import Flask, Blueprint, request, jsonify, render_template, send_file, abort, redirect, Response
-
-from rdkit import Chem
-from rdkit.Chem import AllChem
-from rdkit.Chem.Draw import rdMolDraw2D
-
-from NMR_Prediction.valid import validate_smiles
-
-bp = Blueprint("cascade", __name__)
-
-redis_client = redis.StrictRedis(
-    host="localhost", port=6379, db=0, decode_responses=True
-)
-
-# ── Pages ─────────────────────────────────────────────────────────────────────
-@bp.route("/")
-@bp.route("")
-@bp.route("/predict/")
-@bp.route("/predict")
-@bp.route("/home/")
-@bp.route("/about/")
-def predict():
-    return render_template("cascade/predict.html")
-
-
-# ── Job submission ────────────────────────────────────────────────────────────
-def _submit_job(smiles, type_):
-    if not validate_smiles(smiles):
-        return jsonify({"message": "Input molecule is not allowed", "task_id": None})
-    task_id = uuid.uuid4().hex
-    redis_client.set(f"task_detail_{task_id}",
-                     json.dumps({"smiles": smiles, "type_": type_}))
-    redis_client.rpush("task_queue", task_id)
-    return task_id
-
-
-@bp.route("/predict_NMR_C/", methods=["POST"])
-def predict_NMR_C():
-    result = _submit_job(request.form["smiles"], request.form["type_"])
-    if not isinstance(result, str):
-        return result
-    return jsonify({"message": "Molecule submitted to C queue", "task_id": result})
-
-
-@bp.route("/predict_NMR_H/", methods=["POST"])
-def predict_NMR_H():
-    result = _submit_job(request.form["smiles"], request.form["type_"])
-    if not isinstance(result, str):
-        return result
-    return jsonify({"message": "Molecule submitted to H queue", "task_id": result})
-
-
-# ── check_task ────────────────────────────────────────────────────────────────
-@bp.route("/check_task/")
-def check_task():
-    raw = redis_client.get(f"task_result_{request.args['task_id']}")
-    if not raw:
-        return "running", 200
-    result = json.loads(raw)
-    if "errMessage" in result:
-        return "Error1", 200
-    return "done", 200
-
-
-# ── get_result ────────────────────────────────────────────────────────────────
-@bp.route("/get_result/")
-def get_result():
-    task_id = request.args["task_id"]
-    raw = redis_client.get(f"task_result_{task_id}")
-    if not raw:
-        return jsonify({"error": "Result not found"}), 404
-
-    result = json.loads(raw)
-    if "errMessage" in result:
-        return jsonify({"error": result["errMessage"]}), 500
-
-    smiles  = result["smiles"]
-    nucleus = result.get("type_", "C")  # "C" or "H"
-
-    weighted_shift_txt = result["weightedShiftTxt"]
-    shift_map = {}
-    for item in filter(None, weighted_shift_txt.split(";")):
-        parts = item.split(",")
-        if len(parts) == 2:
-            shift_map[int(parts[0])] = parts[1]
-
-    if nucleus == "H":
-        # ¹H: draw with explicit H so H atoms are visible
-        mol = Chem.MolFromSmiles(smiles)
-        mol = Chem.AddHs(mol)
-        AllChem.Compute2DCoords(mol)
-        mol_draw = rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=True)
-        n_label  = mol.GetNumAtoms()
-        drawer   = rdMolDraw2D.MolDraw2DSVG(700, 500)
-    else:
-        # ¹³C: skeletal structure, no explicit H
-        mol = Chem.MolFromSmiles(smiles)
-        AllChem.Compute2DCoords(mol)
-        mol_draw = rdMolDraw2D.PrepareMolForDrawing(mol, kekulize=True)
-        n_label  = mol.GetNumAtoms()
-        drawer   = rdMolDraw2D.MolDraw2DSVG(600, 450)
-
-    opts = drawer.drawOptions()
-    for atom_1idx, shift_val in shift_map.items():
-        atom_0idx = atom_1idx - 1
-        if atom_0idx < n_label:
-            opts.atomLabels[atom_0idx] = shift_val
-    opts.clearBackground = False
-    opts.bondLineWidth = 1
-    opts.padding = 0.15
-    opts.additionalAtomLabelPadding = 0.1
-
-    drawer.DrawMolecule(mol_draw)
-    drawer.FinishDrawing()
-    svg = drawer.GetDrawingText().replace("svg:", "").replace(":svg", "")
-
-    return jsonify({
-        "svg": svg,
-        "smiles": smiles,
-        "nucleus": nucleus,
-        "conf_sdfs": result.get("conf_sdfs", []),
-        "weightedShift": weighted_shift_txt,
-        "confShift": result["confShiftTxt"],
-        "relative_E": result["relative_E"],
-        "taskId": task_id,
-    })
-
-
-# ── Download as CSV ───────────────────────────────────────────────────────────
-@bp.route("/download/<task_id>/")
-def download(task_id):
-    raw = redis_client.get(f"task_result_{task_id}")
-    if not raw:
-        abort(404)
-    result = json.loads(raw)
-    if "errMessage" in result:
-        abort(404)
-
-    nucleus = result.get("type_", "C")
-    header = f"Atom Index,Predicted {'1H' if nucleus == 'H' else '13C'} Shift (ppm)"
-    lines = [header]
-    for item in filter(None, result["weightedShiftTxt"].split(";")):
-        parts = item.split(",")
-        if len(parts) == 2:
-            lines.append(f"{parts[0]},{parts[1]}")
-
-    return Response(
-        "\n".join(lines),
-        mimetype="text/csv",
         headers={"Content-Disposition": f"attachment; filename=cascade_{task_id[:8]}.csv"}
     )
 
@@ -170,6 +11,13 @@ def create_app():
     def root():
         return redirect("/cascade_v1/predict/")
 
+    # Allow HF Spaces to embed the app in an iframe
+    @app.after_request
+    def remove_iframe_restriction(response):
+        response.headers.pop("X-Frame-Options", None)
+        response.headers["Content-Security-Policy"] = "frame-ancestors *"
+        return response
+
     return app
 
 app = create_app()