Fix missing Kornia import in GHOST 2.0 blender

ghost_runtime.py

@@ -30,9 +30,35 @@ _LOCK = threading.Lock()
 _ENGINE: "Ghost2Engine | None" = None
 
 
+def _patch_source(root: Path) -> None:
+    """Apply small inference fixes that are missing from the upstream release."""
+    crops = root / "src" / "utils" / "crops.py"
+    text = crops.read_text()
+    text = text.replace(
+        "from repos.emoca.gdl.datasets.ImageDatasetHelpers import bbox2point\n", ""
+    )
+    crops.write_text(text)
+
+    embedder = root / "src" / "aligner" / "embedder.py"
+    text = embedder.read_text()
+    text = text.replace("weights='DEFAULT'", "weights=None")
+    embedder.write_text(text)
+
+    # Upstream BlenderGenerator calls kornia_morphology but relies on an import
+    # in the training module. Direct inference imports the generator itself, so
+    # make the dependency explicit in that module.
+    generator = root / "src" / "blender" / "generator.py"
+    text = generator.read_text()
+    import_line = "import src.utils.kornia_morphology as kornia_morphology\n"
+    if import_line not in text:
+        text = text.replace("import torch.nn.functional as F\n", f"import torch.nn.functional as F\n{import_line}")
+    generator.write_text(text)
+
+
 def _prepare_source() -> Path:
     marker = SOURCE_DIR / ".ready"
     if marker.exists():
+        _patch_source(SOURCE_DIR)
         return SOURCE_DIR
     shutil.rmtree(SOURCE_DIR, ignore_errors=True)
     SOURCE_DIR.mkdir(parents=True, exist_ok=True)
@@ -49,18 +75,8 @@ def _prepare_source() -> Path:
             with archive.open(member) as source, destination.open("wb") as target:
                 shutil.copyfileobj(source, target)
 
-    # emoca_crop is unused during inference; removing this optional training import
-    # avoids pulling DECA/EMOCA/PyTorch3D into the ZeroGPU image.
-    crops = SOURCE_DIR / "src" / "utils" / "crops.py"
-    text = crops.read_text()
-    text = text.replace(
-        "from repos.emoca.gdl.datasets.ImageDatasetHelpers import bbox2point\n", ""
-    )
-    crops.write_text(text)
-    embedder = SOURCE_DIR / "src" / "aligner" / "embedder.py"
-    text = embedder.read_text()
-    text = text.replace("weights='DEFAULT'", "weights=None")
-    embedder.write_text(text)
+    # Remove optional training-only imports and patch the upstream inference path.
+    _patch_source(SOURCE_DIR)
     marker.write_text("ready")
     return SOURCE_DIR