Spaces:

monurcan
/

simit

Running on Zero

simit / demo_models.py

Recreate Triton's driver in the ZeroGPU worker (kernels were compiled for the wrong GPU)

c9eba3e verified 26 days ago

12 kB

"""Models of the demo and how one request runs on ZeroGPU.
ZeroGPU forks the Space process for every ``@spaces.GPU`` call and kills the
fork afterwards. So the weights are loaded once, on CPU, in the main process
(no CUDA there), and each call moves only the chosen model to the GPU and
builds a fresh SIMIT engine around it.
"""

from __future__ import annotations

import contextlib
import json
import os
import queue
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import torch

from simit import SIMIT, SIMITConfig

HERE = Path(__file__).parent
PRESETS = json.loads((HERE / "presets.json").read_text())
BUDGETS = [20, 40, 60, 90, 180]           # seconds of GPU time per request
ANSWER_TOKENS = 128
ANSWER_RESERVE = 5.0                      # seconds kept for the final SIMIT-ICL answer


@dataclass
class ModelSpec:
    key: str
    label: str
    repo: str
    kind: str                               # "bagel" | "lance" | "hf"
    gpu_size: str                           # ZeroGPU size: "large" (48 GB) or "xlarge" (96 GB)
    image_generator: Optional[str] = None
    blurb: str = ""
    _weights: object = None
    _lock: threading.Lock = field(default_factory=threading.Lock)

    # ----------------------------------------------------- main process (CPU)
    def load(self):
        """Load the weights on CPU (main process; nothing touches CUDA)."""
        with self._lock:
            if self._weights is None:
                t = time.time()
                self._weights = _LOADERS[self.kind](self)
                print(f"[demo] loaded {self.label} on CPU in {time.time() - t:.0f}s", flush=True)
        return self._weights

    # ------------------------------------------------- inside @spaces.GPU call
    def build(self, preset: dict) -> SIMIT:
        """Move the weights to the GPU and wrap them in a SIMIT engine."""
        w = self.load()
        # Without ZeroGPU every request runs in this process: keep one model on the GPU at a time.
        # (On ZeroGPU each request is a fresh fork in which nothing is on the GPU yet.)
        if _ON_GPU[0] not in (None, self.key):
            MODELS[_ON_GPU[0]]._to("cpu")
            torch.cuda.empty_cache()
        _ON_GPU[0] = self.key
        cfg = SIMITConfig(**{k: v for k, v in preset.items() if k in SIMITConfig.__dataclass_fields__})
        if self.kind == "bagel":
            from simit.backends.bagel import BagelBackend
            backend = BagelBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
            return SIMIT(backend, config=cfg, model_id=self.repo)
        if self.kind == "lance":
            from simit.backends.lance import LanceBackend
            backend = LanceBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
            return SIMIT(backend, config=cfg, model_id=self.repo)
        from simit.backends.hf import HFBackend
        from simit.backends.imagegen import DiffusersImageGenerator
        model, processor, pipe = w
        model.to("cuda")
        pipe.to("cuda")
        backend = HFBackend(model, processor=processor, max_batch_size=8, max_batch_tokens=32768)
        return SIMIT(backend, image_generator=DiffusersImageGenerator(pipe), config=cfg, model_id=self.repo)

    def _to(self, device: str):
        if self.kind in ("bagel", "lance"):
            self._weights.to(device)
        else:
            model, _, pipe = self._weights
            model.to(device)
            pipe.to(device)

    def preset(self, budget: int) -> dict:
        return dict(PRESETS[self.key][str(budget)])


_ON_GPU = [None]   # key of the model whose weights are on the GPU in this process


@contextlib.contextmanager
def _fp8_without_gpu():
    """transformers dequantizes FP8 checkpoints to bf16 when it sees no GPU at load
    time. FP8 weights only need a GPU to *run*, so let the load keep them in FP8
    on CPU; they are moved to the GPU inside each request."""
    saved = torch.cuda.is_available, torch.cuda.get_device_capability
    torch.cuda.is_available = lambda: True
    torch.cuda.get_device_capability = lambda *a, **k: (9, 0)
    try:
        yield
    finally:
        torch.cuda.is_available, torch.cuda.get_device_capability = saved


def _load_bagel(spec):
    from simit.backends.bagel import BagelModel
    return BagelModel.load(spec.repo, device="cpu")


def _load_lance(spec):
    from simit.backends.lance import LanceModel
    return LanceModel.load(spec.repo, device="cpu")


def _load_hf(spec):
    from diffusers import DiffusionPipeline
    from transformers import AutoModelForMultimodalLM, AutoProcessor

    from simit.backends.hf import _dense_fp8_fix
    os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1")  # Triton FP8 kernels: no nvcc JIT at runtime
    processor = AutoProcessor.from_pretrained(spec.repo)
    kw = dict(device_map="cpu", dtype=torch.bfloat16)
    _dense_fp8_fix(spec.repo, kw)  # transformers would skip FP8 conversion of gate_proj for this checkpoint
    with _fp8_without_gpu():
        model = AutoModelForMultimodalLM.from_pretrained(spec.repo, **kw)
    model.eval()
    pipe = DiffusionPipeline.from_pretrained(spec.image_generator, dtype=torch.bfloat16)
    return model, processor, pipe


_LOADERS = {"bagel": _load_bagel, "lance": _load_lance, "hf": _load_hf}

MODELS = {m.key: m for m in [
    ModelSpec("bagel", "BAGEL-7B-MoT", "ByteDance-Seed/BAGEL-7B-MoT", "bagel", "large",
              blurb="Unified model: draws its own photos and writes code for charts and diagrams."),
    ModelSpec("lance", "Lance", "bytedance-research/Lance", "lance", "large",
              blurb="Small unified model: fast native image generation."),
    ModelSpec("qwen", "Qwen3.8-27B (FP8) + FLUX.2-klein", "Qwen/Qwen3.8-27B-FP8", "hf", "xlarge",
              image_generator="black-forest-labs/FLUX.2-klein-4B",
              blurb="Standard VLM: FLUX.2-klein draws the photos; it writes code for charts and diagrams."),
]}
ENABLED = [k for k in os.environ.get("SIMIT_DEMO_MODELS", "bagel,lance,qwen").split(",") if k in MODELS]


# ---------------------------------------------------------------- a request
# Pipeline call tags (``on_progress``) -> what the loading screen shows.
STAGE_OF_TAG = {"synthesis": "synthesizing", "synthesis_fresh": "synthesizing", "revise": "synthesizing",
                "route": "routing", "skill_spec": "coding", "render": "coding", "image": "drawing",
                "critic": "checking", "critic_score": "checking"}


def _refresh_triton_driver():
    """ZeroGPU forks this GPU worker from a main process whose torch.cuda queries are faked.
    Triton's driver copies torch.cuda.get_device_capability / current_device when it is first
    created, so a driver created in the main process (e.g. by fla's import-time device probe)
    keeps reporting the fake GPU, and kernels get compiled for the wrong architecture ("no kernel
    image is available for execution on the device"). Recreate it here, with the real GPU."""
    if not torch.cuda.is_available():
        return
    try:
        from triton.backends.nvidia.driver import CudaDriver
        from triton.runtime import driver
    except Exception:
        return
    driver.set_active(CudaDriver())


def run(spec: ModelSpec, image, question: str, budget: int, always: bool):
    """Generator of progress events for one request (runs inside the GPU call):
    ("stage", name, info) / ("greedy", ZeroShot, k) / ("demo", Demo) / ("final", answer, info)."""
    t_start = time.time()
    preset = spec.preset(budget)
    yield ("stage", "warmup", {"elapsed": 0.0})
    _refresh_triton_driver()
    sim = spec.build(preset)
    try:
        yield ("stage", "solving", {"elapsed": time.time() - t_start})
        zs = sim.zero_shot(image, question, max_new_tokens=ANSWER_TOKENS)   # the greedy answer + p0 (ABA input)
        k_budget = sim.config.budget(zs.confidence)
        yield ("greedy", zs, k_budget)

        # what is left of the budget, keeping room for the final answer (measured: ~1-1.5 s after the last demo)
        time_left = budget - (time.time() - t_start) - ANSWER_RESERVE
        if (k_budget == 0 and not always) or time_left < 3:
            if time_left < 3:
                reason = "no_time"
            else:
                reason = "confident" if preset.get("imagine_by_default", True) else "policy"
            yield ("final", zs.answer, {"demos": 0, "skipped": reason, "elapsed": time.time() - t_start})
            return

        events: "queue.Queue" = queue.Queue()
        current = ["synthesis"]
        result = {}

        def work():
            try:
                result["det"] = sim.imagine(image, question, k=preset["k_max"] if always else None,
                                            return_details=True, max_new_tokens=ANSWER_TOKENS, time_limit=time_left,
                                            on_demo=events.put, on_progress=lambda tag: current.__setitem__(0, tag))
            except Exception as e:  # surfaced below
                result["error"] = e
        th = threading.Thread(target=work, daemon=True)
        th.start()
        while th.is_alive() or not events.empty():
            try:
                yield ("demo", events.get(timeout=0.4))
                continue
            except queue.Empty:
                pass
            stage = STAGE_OF_TAG.get(current[0].split(":")[0], "synthesizing")
            yield ("stage", stage, {"elapsed": time.time() - t_start, "limit": budget})
        if "error" in result:
            raise result["error"]
        det = result["det"]
        if not det.demos:
            yield ("final", zs.answer, {"demos": 0, "skipped": "none_passed", "elapsed": time.time() - t_start,
                                        "stats": det.stats})
            return
        yield ("stage", "improving", {"elapsed": time.time() - t_start})
        answer = sim.answer(image, question, det.demos, max_new_tokens=ANSWER_TOKENS)
        yield ("final", answer, {"demos": len(det.demos), "elapsed": time.time() - t_start, "stats": det.stats})
    finally:
        sim.close()


def code_snippet(spec: ModelSpec, budget: int, question: str, always: bool = False) -> str:
    """Copy-paste code reproducing this request with the library."""
    p = spec.preset(budget)
    force = always or not p.get("imagine_by_default", True)
    keys = ["k_max", "epsilon", "A0", "A1", "B", "t_low", "t_high", "verify", "verify_think", "use_skills",
            "attempts_per_slot", "verify_rounds"]
    args = [f"{k}={round(p[k], 3) if isinstance(p[k], float) else p[k]!r}" for k in keys if k in p]
    lines, line = [], ""
    for a in args:  # wrap the config arguments, a few per line
        if line and len(line) + len(a) > 62:
            lines.append(line.rstrip())
            line = ""
        line += a + ", "
    lines.append(line.rstrip(", "))
    cfg = "\n        " + "\n        ".join(lines) + ",\n    "
    gen = f',\n    image_generator="{spec.image_generator}",' if spec.image_generator else ","
    q = question.strip().replace('"""', "'''") or "What is shown in this image?"
    steps = f"\n    image_steps={p['image_steps']}," if p.get("image_steps") else ""
    return f'''# pip install simit
from simit import SIMIT, SIMITConfig
model = SIMIT.from_pretrained(
    "{spec.repo}"{gen}
    config=SIMITConfig({cfg}),{steps}
)
image = "my_image.jpg"
question = """{q}"""
greedy = model.greedy(image, question)                          # base model, zero-shot
demos = model.imagine(image, question{f", k={p['k_max']}" if force else ""}, time_limit={max(5, budget - 15)})   # imagined examples
improved = model.answer(image, question, demos)                 # SIMIT-ICL
print(greedy, "->", improved)
from simit import save_demos
save_demos(demos, "imagined/")                                  # PNGs + jsonl, e.g. to fine-tune on
'''

	"""Models of the demo and how one request runs on ZeroGPU.

	ZeroGPU forks the Space process for every ``@spaces.GPU`` call and kills the
	fork afterwards. So the weights are loaded once, on CPU, in the main process
	(no CUDA there), and each call moves only the chosen model to the GPU and
	builds a fresh SIMIT engine around it.
	"""

	from __future__ import annotations

	import contextlib
	import json
	import os
	import queue
	import threading
	import time
	from dataclasses import dataclass, field
	from pathlib import Path
	from typing import Optional

	import torch

	from simit import SIMIT, SIMITConfig

	HERE = Path(__file__).parent
	PRESETS = json.loads((HERE / "presets.json").read_text())
	BUDGETS = [20, 40, 60, 90, 180] # seconds of GPU time per request
	ANSWER_TOKENS = 128
	ANSWER_RESERVE = 5.0 # seconds kept for the final SIMIT-ICL answer


	@dataclass
	class ModelSpec:
	key: str
	label: str
	repo: str
	kind: str # "bagel" \| "lance" \| "hf"
	gpu_size: str # ZeroGPU size: "large" (48 GB) or "xlarge" (96 GB)
	image_generator: Optional[str] = None
	blurb: str = ""
	_weights: object = None
	_lock: threading.Lock = field(default_factory=threading.Lock)

	# ----------------------------------------------------- main process (CPU)
	def load(self):
	"""Load the weights on CPU (main process; nothing touches CUDA)."""
	with self._lock:
	if self._weights is None:
	t = time.time()
	self._weights = _LOADERS[self.kind](self)
	print(f"[demo] loaded {self.label} on CPU in {time.time() - t:.0f}s", flush=True)
	return self._weights

	# ------------------------------------------------- inside @spaces.GPU call
	def build(self, preset: dict) -> SIMIT:
	"""Move the weights to the GPU and wrap them in a SIMIT engine."""
	w = self.load()
	# Without ZeroGPU every request runs in this process: keep one model on the GPU at a time.
	# (On ZeroGPU each request is a fresh fork in which nothing is on the GPU yet.)
	if _ON_GPU[0] not in (None, self.key):
	MODELS[_ON_GPU[0]]._to("cpu")
	torch.cuda.empty_cache()
	_ON_GPU[0] = self.key
	cfg = SIMITConfig(**{k: v for k, v in preset.items() if k in SIMITConfig.__dataclass_fields__})
	if self.kind == "bagel":
	from simit.backends.bagel import BagelBackend
	backend = BagelBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
	return SIMIT(backend, config=cfg, model_id=self.repo)
	if self.kind == "lance":
	from simit.backends.lance import LanceBackend
	backend = LanceBackend(w, device="cuda", image_steps=preset.get("image_steps"), use_cuda_graphs=False)
	return SIMIT(backend, config=cfg, model_id=self.repo)
	from simit.backends.hf import HFBackend
	from simit.backends.imagegen import DiffusersImageGenerator
	model, processor, pipe = w
	model.to("cuda")
	pipe.to("cuda")
	backend = HFBackend(model, processor=processor, max_batch_size=8, max_batch_tokens=32768)
	return SIMIT(backend, image_generator=DiffusersImageGenerator(pipe), config=cfg, model_id=self.repo)

	def _to(self, device: str):
	if self.kind in ("bagel", "lance"):
	self._weights.to(device)
	else:
	model, _, pipe = self._weights
	model.to(device)
	pipe.to(device)

	def preset(self, budget: int) -> dict:
	return dict(PRESETS[self.key][str(budget)])


	_ON_GPU = [None] # key of the model whose weights are on the GPU in this process


	@contextlib.contextmanager
	def _fp8_without_gpu():
	"""transformers dequantizes FP8 checkpoints to bf16 when it sees no GPU at load
	time. FP8 weights only need a GPU to run, so let the load keep them in FP8
	on CPU; they are moved to the GPU inside each request."""
	saved = torch.cuda.is_available, torch.cuda.get_device_capability
	torch.cuda.is_available = lambda: True
	torch.cuda.get_device_capability = lambda a, *k: (9, 0)
	try:
	yield
	finally:
	torch.cuda.is_available, torch.cuda.get_device_capability = saved


	def _load_bagel(spec):
	from simit.backends.bagel import BagelModel
	return BagelModel.load(spec.repo, device="cpu")


	def _load_lance(spec):
	from simit.backends.lance import LanceModel
	return LanceModel.load(spec.repo, device="cpu")


	def _load_hf(spec):
	from diffusers import DiffusionPipeline
	from transformers import AutoModelForMultimodalLM, AutoProcessor

	from simit.backends.hf import _dense_fp8_fix
	os.environ.setdefault("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "1") # Triton FP8 kernels: no nvcc JIT at runtime
	processor = AutoProcessor.from_pretrained(spec.repo)
	kw = dict(device_map="cpu", dtype=torch.bfloat16)
	_dense_fp8_fix(spec.repo, kw) # transformers would skip FP8 conversion of gate_proj for this checkpoint
	with _fp8_without_gpu():
	model = AutoModelForMultimodalLM.from_pretrained(spec.repo, **kw)
	model.eval()
	pipe = DiffusionPipeline.from_pretrained(spec.image_generator, dtype=torch.bfloat16)
	return model, processor, pipe


	_LOADERS = {"bagel": _load_bagel, "lance": _load_lance, "hf": _load_hf}

	MODELS = {m.key: m for m in [
	ModelSpec("bagel", "BAGEL-7B-MoT", "ByteDance-Seed/BAGEL-7B-MoT", "bagel", "large",
	blurb="Unified model: draws its own photos and writes code for charts and diagrams."),
	ModelSpec("lance", "Lance", "bytedance-research/Lance", "lance", "large",
	blurb="Small unified model: fast native image generation."),
	ModelSpec("qwen", "Qwen3.8-27B (FP8) + FLUX.2-klein", "Qwen/Qwen3.8-27B-FP8", "hf", "xlarge",
	image_generator="black-forest-labs/FLUX.2-klein-4B",
	blurb="Standard VLM: FLUX.2-klein draws the photos; it writes code for charts and diagrams."),
	]}
	ENABLED = [k for k in os.environ.get("SIMIT_DEMO_MODELS", "bagel,lance,qwen").split(",") if k in MODELS]


	# ---------------------------------------------------------------- a request
	# Pipeline call tags (``on_progress``) -> what the loading screen shows.
	STAGE_OF_TAG = {"synthesis": "synthesizing", "synthesis_fresh": "synthesizing", "revise": "synthesizing",
	"route": "routing", "skill_spec": "coding", "render": "coding", "image": "drawing",
	"critic": "checking", "critic_score": "checking"}


	def _refresh_triton_driver():
	"""ZeroGPU forks this GPU worker from a main process whose torch.cuda queries are faked.
	Triton's driver copies torch.cuda.get_device_capability / current_device when it is first
	created, so a driver created in the main process (e.g. by fla's import-time device probe)
	keeps reporting the fake GPU, and kernels get compiled for the wrong architecture ("no kernel
	image is available for execution on the device"). Recreate it here, with the real GPU."""
	if not torch.cuda.is_available():
	return
	try:
	from triton.backends.nvidia.driver import CudaDriver
	from triton.runtime import driver
	except Exception:
	return
	driver.set_active(CudaDriver())


	def run(spec: ModelSpec, image, question: str, budget: int, always: bool):
	"""Generator of progress events for one request (runs inside the GPU call):
	("stage", name, info) / ("greedy", ZeroShot, k) / ("demo", Demo) / ("final", answer, info)."""
	t_start = time.time()
	preset = spec.preset(budget)
	yield ("stage", "warmup", {"elapsed": 0.0})
	_refresh_triton_driver()
	sim = spec.build(preset)
	try:
	yield ("stage", "solving", {"elapsed": time.time() - t_start})
	zs = sim.zero_shot(image, question, max_new_tokens=ANSWER_TOKENS) # the greedy answer + p0 (ABA input)
	k_budget = sim.config.budget(zs.confidence)
	yield ("greedy", zs, k_budget)

	# what is left of the budget, keeping room for the final answer (measured: ~1-1.5 s after the last demo)
	time_left = budget - (time.time() - t_start) - ANSWER_RESERVE
	if (k_budget == 0 and not always) or time_left < 3:
	if time_left < 3:
	reason = "no_time"
	else:
	reason = "confident" if preset.get("imagine_by_default", True) else "policy"
	yield ("final", zs.answer, {"demos": 0, "skipped": reason, "elapsed": time.time() - t_start})
	return

	events: "queue.Queue" = queue.Queue()
	current = ["synthesis"]
	result = {}

	def work():
	try:
	result["det"] = sim.imagine(image, question, k=preset["k_max"] if always else None,
	return_details=True, max_new_tokens=ANSWER_TOKENS, time_limit=time_left,
	on_demo=events.put, on_progress=lambda tag: current.__setitem__(0, tag))
	except Exception as e: # surfaced below
	result["error"] = e
	th = threading.Thread(target=work, daemon=True)
	th.start()
	while th.is_alive() or not events.empty():
	try:
	yield ("demo", events.get(timeout=0.4))
	continue
	except queue.Empty:
	pass
	stage = STAGE_OF_TAG.get(current[0].split(":")[0], "synthesizing")
	yield ("stage", stage, {"elapsed": time.time() - t_start, "limit": budget})
	if "error" in result:
	raise result["error"]
	det = result["det"]
	if not det.demos:
	yield ("final", zs.answer, {"demos": 0, "skipped": "none_passed", "elapsed": time.time() - t_start,
	"stats": det.stats})
	return
	yield ("stage", "improving", {"elapsed": time.time() - t_start})
	answer = sim.answer(image, question, det.demos, max_new_tokens=ANSWER_TOKENS)
	yield ("final", answer, {"demos": len(det.demos), "elapsed": time.time() - t_start, "stats": det.stats})
	finally:
	sim.close()


	def code_snippet(spec: ModelSpec, budget: int, question: str, always: bool = False) -> str:
	"""Copy-paste code reproducing this request with the library."""
	p = spec.preset(budget)
	force = always or not p.get("imagine_by_default", True)
	keys = ["k_max", "epsilon", "A0", "A1", "B", "t_low", "t_high", "verify", "verify_think", "use_skills",
	"attempts_per_slot", "verify_rounds"]
	args = [f"{k}={round(p[k], 3) if isinstance(p[k], float) else p[k]!r}" for k in keys if k in p]
	lines, line = [], ""
	for a in args: # wrap the config arguments, a few per line
	if line and len(line) + len(a) > 62:
	lines.append(line.rstrip())
	line = ""
	line += a + ", "
	lines.append(line.rstrip(", "))
	cfg = "\n " + "\n ".join(lines) + ",\n "
	gen = f',\n image_generator="{spec.image_generator}",' if spec.image_generator else ","
	q = question.strip().replace('"""', "'''") or "What is shown in this image?"
	steps = f"\n image_steps={p['image_steps']}," if p.get("image_steps") else ""
	return f'''# pip install simit
	from simit import SIMIT, SIMITConfig

	model = SIMIT.from_pretrained(
	"{spec.repo}"{gen}
	config=SIMITConfig({cfg}),{steps}
	)

	image = "my_image.jpg"
	question = """{q}"""

	greedy = model.greedy(image, question) # base model, zero-shot
	demos = model.imagine(image, question{f", k={p['k_max']}" if force else ""}, time_limit={max(5, budget - 15)}) # imagined examples
	improved = model.answer(image, question, demos) # SIMIT-ICL
	print(greedy, "->", improved)

	from simit import save_demos
	save_demos(demos, "imagined/") # PNGs + jsonl, e.g. to fine-tune on
	'''