File size: 6,490 Bytes
afff449
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a42bc3b
 
afff449
 
 
 
 
 
a42bc3b
 
 
 
 
 
 
 
 
afff449
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a42bc3b
 
 
 
 
 
 
afff449
 
 
 
a42bc3b
 
 
afff449
 
 
a42bc3b
afff449
 
 
 
a42bc3b
 
 
 
 
afff449
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""
Compute provider selection (Wave A β€” Batch 6 / HP-1 + HP-2).

Resolves which ComputeProvider serves generation, per
``config.HOMEPILOT_COMPUTE_MODE``:

  * ``local``            β†’ LocalComputeProvider (today's behaviour; the default)
  * ``ollabridge_cloud`` β†’ OllaBridgeCloudComputeProvider
  * ``auto``             β†’ local GPU when healthy, else a linked OllaBridge
                           device, else local (offline β€” status says so)

``compute_status()`` backs the plain-language status UX (HP-2): normal users see
"Using this PC β€” Private GPU" or an honest "your PC is offline" message, never
endpoint/API-key configuration.
"""

from __future__ import annotations

from typing import Any, Optional

from .base import ComputeProvider, GeneratedMedia
from .local import LocalComputeProvider
from .ollabridge_cloud import OllaBridgeCloudComputeProvider

__all__ = [
    "ComputeProvider",
    "GeneratedMedia",
    "LocalComputeProvider",
    "OllaBridgeCloudComputeProvider",
    "ComputeRouter",
    "route_chat",
    "get_compute_provider",
    "resolve_mode",
    "compute_status",
]


def __getattr__(name: str) -> Any:
    # Lazily expose the router without importing it at package import time
    # (keeps the import graph light for callers that only need status helpers).
    if name in ("ComputeRouter", "route_chat"):
        from .router import ComputeRouter, route_chat
        return {"ComputeRouter": ComputeRouter, "route_chat": route_chat}[name]
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


def _build_cloud() -> Optional[OllaBridgeCloudComputeProvider]:
    from app.config import (
        OLLABRIDGE_CLOUD_IMAGE_MODEL,
        OLLABRIDGE_CLOUD_TIMEOUT,
        OLLABRIDGE_CLOUD_TOKEN,
        OLLABRIDGE_CLOUD_URL,
        OLLABRIDGE_CLOUD_VIDEO_MODEL,
    )
    if not OLLABRIDGE_CLOUD_URL:
        return None
    return OllaBridgeCloudComputeProvider(
        OLLABRIDGE_CLOUD_URL,
        OLLABRIDGE_CLOUD_TOKEN,
        image_model=OLLABRIDGE_CLOUD_IMAGE_MODEL,
        video_model=OLLABRIDGE_CLOUD_VIDEO_MODEL,
        timeout=OLLABRIDGE_CLOUD_TIMEOUT,
    )


def _configured_mode() -> str:
    from app.config import HOMEPILOT_COMPUTE_MODE
    mode = (HOMEPILOT_COMPUTE_MODE or "local").strip().lower()
    return mode if mode in ("local", "ollabridge_cloud", "auto") else "local"


def _cloud_configured() -> bool:
    """A cloud link is usable only when both a URL and a token are set β€”
    without a token the job API would 401."""
    from app.config import OLLABRIDGE_CLOUD_TOKEN, OLLABRIDGE_CLOUD_URL
    return bool(OLLABRIDGE_CLOUD_URL and OLLABRIDGE_CLOUD_TOKEN)


def _burst_allowed() -> bool:
    """Whether `auto` may burst to a cloud GPU when the local one is offline.
    Free for everyone unless the operator gates it to premium (MB6). The local
    GPU is never affected."""
    from app.config import COMPUTE_BURST_REQUIRES_PREMIUM, PREMIUM_COMPUTE_ENABLED
    return (not COMPUTE_BURST_REQUIRES_PREMIUM) or PREMIUM_COMPUTE_ENABLED


async def resolve_mode(modality: str | None = None) -> str:
    """Resolve the effective mode (never returns ``auto``).

    ``modality`` makes ``auto`` decide against the runtime the request actually
    needs (Ollama for chat, ComfyUI for image/video) rather than a single global
    "is the local GPU up" signal.
    """
    configured = _configured_mode()
    if configured in ("local", "ollabridge_cloud"):
        return configured

    # auto: prefer the local runtime for this modality, then β€” if bursting is
    # allowed β€” a *configured* linked OllaBridge device.
    if await LocalComputeProvider().available(modality):
        return "local"
    if _burst_allowed():
        cloud = _build_cloud()
        if cloud is not None and _cloud_configured() and await cloud.available(modality):
            return "ollabridge_cloud"
    return "local"  # offline (or burst-gated) β€” compute_status() explains


async def get_compute_provider(
    mode: str | None = None, modality: str | None = None
) -> ComputeProvider:
    """Return the provider for the given (or resolved) mode and modality."""
    mode = mode or await resolve_mode(modality)
    if mode == "ollabridge_cloud":
        cloud = _build_cloud()
        if cloud is not None:
            return cloud
    return LocalComputeProvider()


async def compute_status() -> dict[str, Any]:
    """Backs the HP-2 status UX β€” plain language, no infrastructure jargon."""
    from app.config import (
        OLLABRIDGE_CLOUD_TOKEN,
        OLLABRIDGE_CLOUD_URL,
        PREMIUM_COMPUTE_ENABLED,
    )

    local_ok = await LocalComputeProvider().available()
    cloud = _build_cloud()
    cloud_configured = bool(OLLABRIDGE_CLOUD_URL and OLLABRIDGE_CLOUD_TOKEN)
    cloud_ok = await cloud.available() if cloud is not None else False
    mode = await resolve_mode()
    configured = _configured_mode()

    # Burst = we're on cloud because the local GPU is offline (auto fallback).
    burst = mode == "ollabridge_cloud" and configured == "auto" and not local_ok
    # Burst-gated = local offline + a reachable cloud we could use, but premium is
    # required and not granted.
    burst_gated = (
        configured == "auto"
        and not local_ok
        and cloud_configured
        and cloud_ok
        and not _burst_allowed()
    )

    if mode == "local" and local_ok:
        label = "Private GPU"
        message = "Compute: Connected β€” Using this PC β€” Mode: Private GPU"
    elif mode == "ollabridge_cloud" and cloud_ok:
        label = "OllaBridge Cloud"
        message = (
            "Your PC is offline β€” running on a cloud GPU (premium)."
            if burst
            else "Compute: Connected β€” Using your paired GPU via OllaBridge Cloud"
        )
    elif burst_gated:
        label = "Offline"
        message = "Your PC is offline β€” upgrade to premium to run on a cloud GPU, or wait for your PC."
    else:
        label = "Offline"
        message = (
            "Your PC is offline β€” wait for it to come back, use the free cloud "
            "queue, or use a premium GPU."
        )

    return {
        "mode": mode,
        "configured_mode": configured,
        "local_gpu_available": local_ok,
        "cloud_configured": cloud_configured,
        "cloud_reachable": cloud_ok,
        "premium": PREMIUM_COMPUTE_ENABLED,
        "burst": burst,
        "burst_gated": burst_gated,
        "label": label,
        "message": message,
    }