"""Modal backend for FrogQuest's GPU work — the off-Space alternative to ZeroGPU. Deploy once: `modal deploy modal_app.py` The HF Space (running on CPU-basic with FROGQUEST_BACKEND=modal) then invokes these via `modal.Cls.from_name("frogquest", "LLM" | "Flux")()..remote(...)` from llm.py / images.py. Both backends share the SAME prompts, config, JSON extractor (gpu_shared) and JSON schemas (schema) so they can never drift. Those two local files are mounted into the image with `add_local_python_source`. GPU config: - LLM (Nemotron Q8 ~4.3GB): a cheap L4 is plenty. - FLUX.2 klein (~23GB): set FROGQUEST_MODAL_FLUX_GPU at DEPLOY time (default "A10G", 24GB -> needs enable_model_cpu_offload(); use "L4" to save cost or "L40S"/"A100-40GB" if it OOMs). NOTE: re-verify Modal API specifics against current docs (Modal 1.0+): App / Image.pip_install extra_index_url / add_local_python_source / @app.cls / @modal.enter / @modal.method / Cls.from_name / Volume.commit. """ import os import modal from gpu_shared import ( CAMPAIGN_SYSTEM_PROMPT, GGUF_FILE, GGUF_REPO, GUIDANCE, INTENT_SYSTEM_PROMPT, MAX_SIDE, MODEL_ID, N_CTX, STEPS, SYSTEM_PROMPT, build_edit_prompt, build_initial_prompt, extract_json, preload_cuda_libs, ) from schema import CAMPAIGN_RESPONSE_SCHEMA, INTENT_SCHEMA, RESPONSE_SCHEMA app = modal.App("frogquest") # Persist the HF weight cache (FLUX ~23GB + GGUF ~4.3GB) across cold starts -> download once. CACHE_DIR = "/cache" vol = modal.Volume.from_name("frogquest-cache", create_if_missing=True) # Prebuilt CUDA wheel for llama.cpp (same index the Space uses). cu124 wheels run on Modal GPUs. _LLAMA_INDEX = "https://abetlen.github.io/llama-cpp-python/whl/cu124" image = ( modal.Image.debian_slim(python_version="3.12") .apt_install("git") # needed to pip-install diffusers from its git URL .pip_install( "git+https://github.com/huggingface/diffusers.git", # Flux2KleinPipeline (>=0.38) "transformers>=4.44.2", "accelerate", # required by enable_model_cpu_offload() "torch", # CUDA runtime libs the prebuilt cu124 llama-cpp wheel links against (libcudart.so.12, # libcublas.so.12, ...). Installed explicitly so the .so files exist on disk; preload_cuda_libs() # then loads them RTLD_GLOBAL before `from llama_cpp import Llama`. (Same set as the Space's # requirements.txt — torch alone does not reliably put libcudart.so.12 on the loader path.) "nvidia-cuda-runtime-cu12", "nvidia-cublas-cu12", "nvidia-cuda-nvrtc-cu12", "pillow", "sentencepiece", # FLUX.2 text-encoder tokenizer deps (belt-and-suspenders) "protobuf", "huggingface-hub", "hf-transfer", ) .pip_install( "llama-cpp-python==0.3.28", extra_index_url=_LLAMA_INDEX, extra_options="--prefer-binary", ) .env({"HF_HOME": CACHE_DIR, "HF_HUB_ENABLE_HF_TRANSFER": "1"}) .add_local_python_source("schema", "gpu_shared") ) # Configurable at deploy time (read on your machine when you run `modal deploy`). FLUX_GPU = os.environ.get("FROGQUEST_MODAL_FLUX_GPU", "A10G") @app.cls(gpu="L4", image=image, volumes={CACHE_DIR: vol}, scaledown_window=300) class LLM: @modal.enter() def _load(self): vol.reload() # pick up weights another container may have already cached import torch # noqa: F401 (loads CUDA libs RTLD_GLOBAL) preload_cuda_libs() from llama_cpp import Llama self.llm = Llama.from_pretrained( repo_id=GGUF_REPO, filename=GGUF_FILE, n_gpu_layers=-1, n_ctx=N_CTX, # Modal GPU has the VRAM for the full 128k verbose=False, ) vol.commit() # persist the freshly downloaded GGUF (no-op if already cached) @modal.method() def generate_quests(self, todos: str, theme: str) -> dict: system = SYSTEM_PROMPT.replace("{theme}", theme) user = f"Theme: {theme}\nMy to-do list / goals:\n{todos.strip()}" out = self.llm.create_chat_completion( messages=[ {"role": "system", "content": system}, {"role": "user", "content": user}, ], response_format={"type": "json_object", "schema": RESPONSE_SCHEMA}, temperature=0.0, max_tokens=4096, ) return extract_json(out["choices"][0]["message"]["content"]) @modal.method() def generate_campaign(self, goal: str, theme: str, snippets: str = "") -> dict: system = CAMPAIGN_SYSTEM_PROMPT.replace("{theme}", theme) user = f"Theme: {theme}\nLong-term goal:\n{goal.strip()}" if (snippets or "").strip(): user += f"\n\nResearch notes:\n{snippets.strip()}" out = self.llm.create_chat_completion( messages=[ {"role": "system", "content": system}, {"role": "user", "content": user}, ], response_format={"type": "json_object", "schema": CAMPAIGN_RESPONSE_SCHEMA}, temperature=0.0, max_tokens=4096, ) return extract_json(out["choices"][0]["message"]["content"]) @modal.method() def route_intent(self, message: str, context: str) -> dict: user = f"Context:\n{context.strip()}\n\nUser message:\n{message.strip()}" out = self.llm.create_chat_completion( messages=[ {"role": "system", "content": INTENT_SYSTEM_PROMPT}, {"role": "user", "content": user}, ], response_format={"type": "json_object", "schema": INTENT_SCHEMA}, temperature=0.0, max_tokens=256, ) parsed = extract_json(out["choices"][0]["message"]["content"]) if not isinstance(parsed, dict) or parsed.get("intent") not in ( "forge", "add_tasks", "mark_done", "mark_couldnt", "unknown", ): return {"intent": "unknown"} return parsed @app.cls(gpu=FLUX_GPU, image=image, volumes={CACHE_DIR: vol}, scaledown_window=300, enable_memory_snapshot=True) class Flux: @modal.enter(snap=True) def _load_cpu(self): # The heavy part: read + deserialize ~23GB from the Volume into CPU RAM. This is captured # in the memory snapshot, so later cold starts RESTORE it instead of redoing the load. # MUST NOT touch CUDA here — snapshots are CPU-only. vol.reload() import torch from diffusers import Flux2KleinPipeline self._torch = torch self.pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16) vol.commit() # persist freshly downloaded FLUX weights (no-op if already cached) @modal.enter(snap=False) def _to_gpu(self): # Runs on every wake: just the fast PCIe copy of the already-loaded weights into VRAM. # Loads the FULL model resident (no offload) -> no per-gen streaming, so generation is fast. self.pipe.to("cuda") def _gen(self, prompt, image, seed): gen = self._torch.Generator("cuda").manual_seed(int(seed)) result = self.pipe( prompt=prompt, image=image, generator=gen, num_inference_steps=STEPS, guidance_scale=GUIDANCE, height=MAX_SIDE, width=MAX_SIDE, ) return result.images[0] @modal.method() def initial(self, user_photo, art_style: str, scene_prompt: str, seed: int): return self._gen(build_initial_prompt(art_style, scene_prompt), [user_photo], seed) @modal.method() def initials(self, user_photo, art_style: str, scene_prompts: list, seed: int): """Batch counterpart of initial() — all of a forge's scenes in one container call.""" return [self._gen(build_initial_prompt(art_style, p), [user_photo], seed) for p in scene_prompts] @modal.method() def edit(self, base_image, edit_instruction: str, art_style: str, seed: int): return self._gen(build_edit_prompt(art_style, edit_instruction), base_image, seed)