Delete raven/onnx_export.py

raven/onnx_export.py

@@ -1,39 +0,0 @@
-from __future__ import annotations
-from pathlib import Path
-import torch
-import torch.nn as nn
-from .model import RavenFusion
-from .foundation import FoundationEncoder
-
-
-class FusionEnsembleWrapper(nn.Module):
-    KEYS=("semantic","camera","spectral","patch","specialist")
-    def __init__(self,models): super().__init__(); self.models=nn.ModuleList(models)
-    def forward(self,semantic,camera,spectral,patch,specialist,quality,anchors,luma):
-        f={"semantic":semantic,"camera":camera,"spectral":spectral,"patch":patch,"specialist":specialist}; outs=[m(f,quality,anchors,luma,0.0,False) for m in self.models]
-        logits=torch.stack([o["logit"] for o in outs],1); branches=torch.stack([o["branch_logits"] for o in outs],1); gates=torch.stack([o["gates"] for o in outs],1)
-        return logits.mean(1),logits.std(1,unbiased=False),branches.mean(1),gates.mean(1)
-
-
-class FoundationWrapper(nn.Module):
-    def __init__(self,enc): super().__init__(); self.enc=enc
-    def forward(self,x): return self.enc(x)
-
-
-def export_fusion(model_path:str,out_path:str):
-    ck=torch.load(model_path,map_location="cpu",weights_only=False); models=[]
-    for st in ck["ensemble_states"]:
-        m=RavenFusion(ck["dims"],len(ck.get("ai_source_map",{})),len(ck.get("real_source_map",{}))); m.load_state_dict(st); m.eval(); models.append(m)
-    w=FusionEnsembleWrapper(models).eval(); d=ck["dims"]; inputs=[torch.randn(1,d[k]) for k in ("semantic","camera","spectral","patch","specialist")]+[torch.randn(1,d["quality"]),torch.randn(1,5),torch.tensor([.3])]
-    Path(out_path).parent.mkdir(parents=True,exist_ok=True)
-    torch.onnx.export(w,tuple(inputs),out_path,input_names=["semantic","camera","spectral","patch","specialist","quality","anchors","luma"],output_names=["logit_mean","logit_std","branch_mean","gates_mean"],dynamic_axes={n:{0:"batch"} for n in ["semantic","camera","spectral","patch","specialist","quality","anchors","luma","logit_mean","logit_std","branch_mean","gates_mean"]},opset_version=18)
-    return out_path
-
-
-def export_foundation(model_path:str,out_path:str):
-    ck=torch.load(model_path,map_location="cpu",weights_only=False); enc=FoundationEncoder(ck["foundation_backend"],allow_fallback=False,fallback_backend=ck["foundation_backend"],pretrained=ck.get("foundation_state") is None).eval();
-    if ck.get("foundation_state") is not None: enc.model.load_state_dict(ck["foundation_state"]["state_dict"])
-    w=FoundationWrapper(enc).eval(); x=torch.randn(1,3,224,224)
-    Path(out_path).parent.mkdir(parents=True,exist_ok=True)
-    torch.onnx.export(w,x,out_path,input_names=["image"],output_names=["embedding"],dynamic_axes={"image":{0:"batch"},"embedding":{0:"batch"}},opset_version=18)
-    return out_path