[feat] LoRA controls and integration for Dreamverse

apps/dreamverse/dreamverse/config.py

@@ -57,18 +57,24 @@ def _resolve_frontend_static_dir_candidates() -> tuple[str, ...]:
         "name": "FastLTX2",
         "model_path": "FastVideo/LTX2-Distilled-Diffusers",
         "config_model_path": "FastVideo/LTX2-Distilled-Diffusers",
+        "lora_repo": "FastVideo/LTX2-OmniNFT-LoRA",
     },
     "fast-ltx23": {
         "name": "FastLTX23",
         "model_path": "FastVideo/LTX-2.3-Distilled-Diffusers",
         "config_model_path": "FastVideo/LTX-2.3-Distilled-Diffusers",
+        "lora_repo": "FastVideo/LTX-2.3-OmniNFT-LoRA",
     },
 }
 
 DEFAULT_MODEL_ID = "fast-ltx2"
 
+ACTIVE_MODEL_ID = (os.getenv("DREAMVERSE_MODEL_ID", "").strip() or DEFAULT_MODEL_ID)
+if ACTIVE_MODEL_ID not in MODEL_REGISTRY:
+    ACTIVE_MODEL_ID = DEFAULT_MODEL_ID
+
 # Active model configuration
-MODEL_CONFIG = MODEL_REGISTRY[DEFAULT_MODEL_ID]
+MODEL_CONFIG = MODEL_REGISTRY[ACTIVE_MODEL_ID]
 
 # Generation limits
 SESSION_TIMEOUT_SECONDS = 300
@@ -166,6 +172,74 @@ def _optional_env(*names: str) -> str | None:
 PROMPT_SAFETY_ENABLED = _env_bool("FASTVIDEO_ENABLE_PROMPT_SAFETY", False)
 DREAMVERSE_MAX_AUTOTUNE = _env_bool("DREAMVERSE_MAX_AUTOTUNE", True)
 
+DREAMVERSE_MODEL_PATH = (os.getenv("DREAMVERSE_MODEL_PATH", "").strip() or None)
+if DREAMVERSE_MODEL_PATH:
+    MODEL_CONFIG = {
+        **MODEL_CONFIG,
+        "model_path": DREAMVERSE_MODEL_PATH,
+        "config_model_path": DREAMVERSE_MODEL_PATH,
+    }
+
+AVAILABLE_LORAS = {
+    "pixar": {
+        "repo": "vrgamedevgirl84/LTX_2.3_Pixar_Toon_Style_LoRa",
+        "trigger": "P1x4r",
+        "model": "fast-ltx23",
+        "position": "prepend",
+        "label": "Pixar Toon",
+    },
+    "transition": {
+        "repo": "valiantcat/LTX-2.3-Transition-LORA",
+        "trigger": "zhuanchang",
+        "model": "fast-ltx23",
+        "position": "append",
+        "label": "Transition",
+    },
+}
+
+
+def _active_model_key() -> str:
+    return ACTIVE_MODEL_ID
+
+
+def _available_styles_for_active_model() -> list[str]:
+    active = _active_model_key()
+    return [k for k, v in AVAILABLE_LORAS.items() if v.get("model") == active]
+
+
+def _resolve_lora_spec(spec: str) -> str | None:
+    spec = (spec or "").strip()
+    if not spec:
+        return None
+    if spec.lower() == "omninft":
+        return MODEL_CONFIG.get("lora_repo")
+    if spec.lower() in AVAILABLE_LORAS:
+        return AVAILABLE_LORAS[spec.lower()]["repo"]
+    return spec
+
+
+def _parse_lora_stack(raw: str) -> list[tuple[str, float]]:
+    out: list[tuple[str, float]] = []
+    for item in (raw or "").split(","):
+        item = item.strip()
+        if not item:
+            continue
+        spec, _, stren = item.partition("@")
+        resolved = _resolve_lora_spec(spec)
+        if resolved:
+            try:
+                strength = float(stren) if stren.strip() else 1.0
+            except ValueError:
+                strength = 1.0
+            out.append((resolved, strength))
+    return out
+
+
+DREAMVERSE_LORA_PATH = _resolve_lora_spec(os.getenv("DREAMVERSE_LORA_PATH", ""))
+DREAMVERSE_LORA_NICKNAME = (os.getenv("DREAMVERSE_LORA_NICKNAME", "omninft").strip() or "omninft")
+DREAMVERSE_LORA_STRENGTH = _env_float("DREAMVERSE_LORA_STRENGTH", 1.0)
+DREAMVERSE_LORA_STACK = _parse_lora_stack(os.getenv("DREAMVERSE_LORA_STACK", ""))
+
 
 def _resolve_devtools_paths(
     default_path: Path,

apps/dreamverse/dreamverse/gpu_pool.py

@@ -33,6 +33,8 @@
     InitAck,
     JoinAck,
     LeaveAck,
+    LoraAck,
+    LoraStackPayload,
     MediaChunk,
     MediaComplete,
     MediaInit,
@@ -80,6 +82,7 @@ class CommandType(Enum):
     USER_STEP = "user_step"
     USER_LEAVE = "user_leave"
     RELOAD_MODEL = "reload_model"
+    APPLY_LORA = "apply_lora"
 
 
 @dataclass
@@ -282,6 +285,25 @@ def _publish(event: StreamEvent) -> None:
                         message=str(e),
                     ))
 
+            elif cmd.type == CommandType.APPLY_LORA:
+                try:
+                    assert isinstance(cmd.payload, LoraStackPayload), (f"APPLY_LORA requires LoraStackPayload, "
+                                                                       f"got {type(cmd.payload).__name__}")
+                    trigger, position = worker.apply_lora_stack(cmd.payload.stack)
+                    response_queue.put(
+                        LoraAck(
+                            user_id=cmd.user_id,
+                            style_trigger=trigger,
+                            style_trigger_position=position,
+                        ))
+                except Exception as e:
+                    print(f"[GPU {gpu_id}] Apply LoRA error: {e}")
+                    traceback.print_exc()
+                    response_queue.put(WorkerError(
+                        user_id=cmd.user_id,
+                        message=str(e),
+                    ))
+
             return True  # Continue loop
 
         if first_cmd is not None:
@@ -359,6 +381,25 @@ def _publish(event: StreamEvent) -> None:
                         message=str(e),
                     ))
 
+            elif cmd.type == CommandType.APPLY_LORA:
+                try:
+                    assert isinstance(cmd.payload, LoraStackPayload), (f"APPLY_LORA requires LoraStackPayload, "
+                                                                       f"got {type(cmd.payload).__name__}")
+                    trigger, position = worker.apply_lora_stack(cmd.payload.stack)
+                    response_queue.put(
+                        LoraAck(
+                            user_id=cmd.user_id,
+                            style_trigger=trigger,
+                            style_trigger_position=position,
+                        ))
+                except Exception as e:
+                    print(f"[GPU {gpu_id}] Apply LoRA error: {e}")
+                    traceback.print_exc()
+                    response_queue.put(WorkerError(
+                        user_id=cmd.user_id,
+                        message=str(e),
+                    ))
+
             elif cmd.type in (CommandType.USER_JOIN, CommandType.USER_STEP, CommandType.USER_LEAVE):
                 event_loop(first_cmd=cmd)
                 break
@@ -729,6 +770,25 @@ async def user_step(
                 raise RuntimeError(f"Unexpected step response for {user_id[:8]}: "
                                    f"{type(response).__name__}")
 
+    async def apply_lora_stack(
+        self,
+        stack: list[tuple[str, float]],
+    ) -> LoraAck:
+        """Re-apply a runtime LoRA stack on this GPU's worker."""
+        self._active = True
+        user_id = "__lora__"
+        payload = LoraStackPayload(stack=stack)
+        response = await self._send_command_tagged(Command(CommandType.APPLY_LORA, payload=payload, user_id=user_id),
+                                                   timeout=120.0)
+        match response:
+            case LoraAck() as ack:
+                return ack
+            case WorkerError(message=msg):
+                raise RuntimeError(f"Apply LoRA failed on GPU {self.gpu_id}: {msg}")
+            case _:
+                raise RuntimeError(f"Unexpected LoRA response on GPU {self.gpu_id}: "
+                                   f"{type(response).__name__}")
+
     async def leave_user(self, user_id: str) -> None:
         """Remove a user from this GPU."""
         try:
@@ -903,6 +963,17 @@ async def _send_queue_updates(self):
                 except Exception:
                     pass
 
+    async def apply_lora_stack(
+        self,
+        stack: list[tuple[str, float]],
+    ) -> dict[int, str | None]:
+        """Re-apply a runtime LoRA stack across all ready GPU workers."""
+        ready_slots = [slot for slot in self.slots.values() if slot.ready]
+        if not ready_slots:
+            raise RuntimeError("No ready GPU workers to apply LoRA stack.")
+        results = await asyncio.gather(*(slot.apply_lora_stack(stack) for slot in ready_slots))
+        return {slot.gpu_id: ack.style_trigger for slot, ack in zip(ready_slots, results, strict=False)}
+
     async def shutdown(self):
         """Shutdown all GPU workers."""
         print("Shutting down GPU pool...")

apps/dreamverse/dreamverse/main.py

@@ -6,18 +6,22 @@
 from contextlib import asynccontextmanager
 from pathlib import Path
 
-from fastapi import FastAPI, WebSocket
+from fastapi import FastAPI, HTTPException, WebSocket
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
 from fastvideo.entrypoints.streaming import build_health_router
 from dreamverse.gpu_pool import GPUPool, get_available_gpus
 from dreamverse.session_logger import SessionEventLogger
 
 from dreamverse.config import (
+    AVAILABLE_LORAS,
     DEVTOOLS_ENABLED,
     FRONTEND_STATIC_DIR_CANDIDATES,
     PROMPT_SAFETY_ENABLED,
     SESSION_LOG_ROOT,
+    _available_styles_for_active_model,
+    _resolve_lora_spec,
 )
 from dreamverse.prompt_enhancer import PromptEnhancer
 from dreamverse.prompt_safety import PromptSafetyFilter
@@ -104,6 +108,83 @@ async def websocket_endpoint(websocket: WebSocket):
     await controller.run()
 
 
+class LoraRequest(BaseModel):
+    strength: float = 1.0
+    styles: dict[str, float] = {}
+    style: str = ""
+
+
+@app.get("/lora/options")
+async def lora_options() -> dict:
+    if not DEVTOOLS_ENABLED:
+        raise HTTPException(status_code=404, detail="Not found")
+    styles = _available_styles_for_active_model()
+    has_base_lora = _resolve_lora_spec("omninft") is not None
+    labels = {"none": "None"}
+    labels.update({k: AVAILABLE_LORAS[k].get("label", k) for k in styles})
+    return {
+        "styles": ["none", *styles],
+        "labels": labels,
+        "has_base_lora": has_base_lora,
+    }
+
+
+@app.post("/lora")
+async def apply_lora(request: LoraRequest) -> dict:
+    if not DEVTOOLS_ENABLED:
+        raise HTTPException(status_code=404, detail="Not found")
+    if runtime.gpu_pool is None:
+        raise HTTPException(status_code=503, detail="GPU pool not ready")
+
+    strength = max(0.0, min(1.0, float(request.strength)))
+    allowed_styles = _available_styles_for_active_model()
+
+    requested = dict(request.styles) if request.styles else {}
+    if not requested and request.style and request.style.strip().lower() != "none":
+        requested = {request.style: 1.0}
+
+    styles_map: dict[str, float] = {}
+    for name, intensity in requested.items():
+        key = str(name).strip().lower()
+        if key in ("", "none"):
+            continue
+        if key not in allowed_styles:
+            raise HTTPException(status_code=400, detail=f"Unknown style for active model: {key}")
+        value = max(0.0, min(1.0, float(intensity)))
+        if value <= 0.0:
+            continue
+        styles_map[key] = value
+
+    stack: list[tuple[str, float]] = []
+    if _resolve_lora_spec("omninft") is not None:
+        stack.append(("omninft", strength))
+    for key, intensity in styles_map.items():
+        stack.append((key, intensity))
+
+    if not stack:
+        raise HTTPException(status_code=400,
+                            detail="Active model has no OmniNFT LoRA and no style selected; nothing to apply.")
+
+    try:
+        triggers = await runtime.gpu_pool.apply_lora_stack(stack)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+    return {
+        "applied": True,
+        "strength": strength,
+        "styles": styles_map,
+        "triggers": {
+            key: AVAILABLE_LORAS[key]["trigger"]
+            for key in styles_map
+        },
+        "gpus": {
+            str(gpu_id): trigger
+            for gpu_id, trigger in triggers.items()
+        },
+    }
+
+
 # Serve an exported frontend bundle when present.
 for static_dir in FRONTEND_STATIC_DIR_CANDIDATES:
     if os.path.isdir(static_dir):

apps/dreamverse/dreamverse/tests/test_entrypoints.py

@@ -171,3 +171,69 @@ def test_mock_server_cli_updates_latency(monkeypatch):
         assert mock_server.LATENCY_MS == 321
     finally:
         mock_server.LATENCY_MS = old_latency_ms
+
+
+def _lora_client(monkeypatch, captured):
+    server_main = _import_server_main(monkeypatch)
+    monkeypatch.setattr(server_main, "DEVTOOLS_ENABLED", True)
+    monkeypatch.setattr(server_main, "_available_styles_for_active_model", lambda: ["pixar", "transition"])
+    monkeypatch.setattr(
+        server_main,
+        "_resolve_lora_spec",
+        lambda spec: "FastVideo/OmniNFT" if str(spec).strip().lower() == "omninft" else spec,
+    )
+
+    class _FakePool:
+
+        async def apply_lora_stack(self, stack):
+            captured["stack"] = list(stack)
+            return {0: "trigger"}
+
+    server_main.runtime.gpu_pool = _FakePool()
+    return TestClient(server_main.app)
+
+
+def test_apply_lora_stacks_multiple_styles_with_intensities(monkeypatch):
+    captured: dict = {}
+    client = _lora_client(monkeypatch, captured)
+    response = client.post("/lora", json={"strength": 0.8, "styles": {"pixar": 0.45, "transition": 0.5}})
+    assert response.status_code == 200
+    body = response.json()
+    assert body["applied"] is True
+    assert body["styles"] == {"pixar": 0.45, "transition": 0.5}
+    assert captured["stack"] == [("omninft", 0.8), ("pixar", 0.45), ("transition", 0.5)]
+
+
+def test_apply_lora_rejects_unknown_style(monkeypatch):
+    captured: dict = {}
+    client = _lora_client(monkeypatch, captured)
+    response = client.post("/lora", json={"strength": 0.5, "styles": {"watercolor": 0.5}})
+    assert response.status_code == 400
+    assert "stack" not in captured
+
+
+def test_apply_lora_accepts_legacy_single_style(monkeypatch):
+    captured: dict = {}
+    client = _lora_client(monkeypatch, captured)
+    response = client.post("/lora", json={"strength": 0.6, "style": "pixar"})
+    assert response.status_code == 200
+    assert captured["stack"] == [("omninft", 0.6), ("pixar", 1.0)]
+
+
+def test_apply_lora_clamps_and_drops_zero_intensity(monkeypatch):
+    captured: dict = {}
+    client = _lora_client(monkeypatch, captured)
+    response = client.post("/lora", json={"strength": 1.5, "styles": {"pixar": 2.0, "transition": 0.0}})
+    assert response.status_code == 200
+    body = response.json()
+    assert body["strength"] == 1.0
+    assert body["styles"] == {"pixar": 1.0}
+    assert captured["stack"] == [("omninft", 1.0), ("pixar", 1.0)]
+
+
+def test_lora_endpoints_hidden_without_devtools(monkeypatch):
+    server_main = _import_server_main(monkeypatch)
+    monkeypatch.setattr(server_main, "DEVTOOLS_ENABLED", False)
+    client = TestClient(server_main.app)
+    assert client.get("/lora/options").status_code == 404
+    assert client.post("/lora", json={"strength": 0.5}).status_code == 404