build-small-hackathon/Hackathon-IA-VisualNovel

1"""Modal backend — LLM + image generation on cloud GPUs, single app deployment.2 3Deploy once:4    modal deploy modal_app.py5 6Download models to volumes before first run:7    modal run modal_app.py::download_model8    modal run modal_app.py::download_image_model9 10Then launch locally:11    uv run python app.py12"""13 14from __future__ import annotations15 16import modal17 18hf_secret = modal.Secret.from_name("huggingface")19 20# =========================================================================== #21#  Container images22# =========================================================================== #23 24llm_image = (25    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12")26    .pip_install(27        "llama-cpp-python",28        extra_options="--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124",29    )30    .pip_install("huggingface_hub")31)32 33painter_image = (34    modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.12")35    .pip_install(36        "torch==2.5.1",37        "torchvision==0.20.1",38        extra_options="--index-url https://download.pytorch.org/whl/cu124",39    )40    .pip_install(41        "diffusers==0.35.1",42        "transformers==4.47.1",43        "accelerate==1.8.1",44        "peft",  # required by diffusers load_lora_weights / fuse_lora45        "safetensors",46        "Pillow",47        "huggingface_hub",48        "rembg",  # background removal for sprites49        "onnxruntime",  # rembg runtime dep50    )51)52 53# =========================================================================== #54#  Volumes  (model weights persist across cold starts)55# =========================================================================== #56 57llm_volume = modal.Volume.from_name("vn-models", create_if_missing=True)58LLM_DIR = "/models"59 60image_volume = modal.Volume.from_name("vn-image-models", create_if_missing=True)61IMAGE_DIR = "/image-models"62 63# =========================================================================== #64#  Single app — both classes deploy together with `modal deploy modal_app.py`65# =========================================================================== #66 67app = modal.App("vn-app")68 69GGUF_REPO = "Qwen/Qwen3-14B-GGUF"70GGUF_FILE = "Qwen3-14B-Q8_0.gguf"71# Painter: SDXL-base-1.0 + ByteDance Lightning LoRA (matches local SdxlLightningPainter)72SDXL_BASE_REPO = "stabilityai/stable-diffusion-xl-base-1.0"73SDXL_BASE_LOCAL = "sdxl-base"74LIGHTNING_LORA_REPO = "ByteDance/SDXL-Lightning"75LIGHTNING_LORA_FILE = "sdxl_lightning_4step_lora.safetensors"76 77 78# --------------------------------------------------------------------------- #79#  Download helpers (run once each)80# --------------------------------------------------------------------------- #81 82 83@app.function(image=llm_image, volumes={LLM_DIR: llm_volume}, timeout=600, secrets=[hf_secret])84def download_model() -> None:85    from huggingface_hub import hf_hub_download86 87    print(f"Downloading {GGUF_REPO}/{GGUF_FILE} ...")88    path = hf_hub_download(repo_id=GGUF_REPO, filename=GGUF_FILE, local_dir=LLM_DIR)89    llm_volume.commit()90    print(f"Saved -> {path}")91 92 93@app.function(94    image=painter_image, volumes={IMAGE_DIR: image_volume}, timeout=900, secrets=[hf_secret]95)96def download_image_model() -> None:97    """Download SDXL-base-1.0 weights (~6.5 GB) to the volume.98 99    The Lightning LoRA (ByteDance/SDXL-Lightning) is small (~400 MB) and loaded100    at inference time directly from HF Hub — no need to pre-download it.101    """102    from huggingface_hub import snapshot_download103 104    print(f"Downloading {SDXL_BASE_REPO} -> {IMAGE_DIR}/{SDXL_BASE_LOCAL} ...")105    snapshot_download(SDXL_BASE_REPO, local_dir=f"{IMAGE_DIR}/{SDXL_BASE_LOCAL}")106    image_volume.commit()107    print("SDXL-base saved to volume.")108 109 110# --------------------------------------------------------------------------- #111#  LLM backend  (llama.cpp + Qwen3-14B on A10G)112# --------------------------------------------------------------------------- #113 114 115@app.cls(116    image=llm_image,117    gpu="A10G",118    volumes={LLM_DIR: llm_volume},119    timeout=300,120    scaledown_window=600,  # keep warm 10 min — reloading the 15 GB GGUF mid-session is worse121)122class ModalLLMBackend:123    @modal.enter()124    def load(self) -> None:125        from llama_cpp import Llama126 127        self.llm = Llama(128            model_path=f"{LLM_DIR}/{GGUF_FILE}",129            n_ctx=8192,130            n_gpu_layers=-1,131            verbose=False,132        )133        print("[modal] LLM loaded on GPU")134 135    @modal.method()136    def complete(self, messages: list[dict], **kw) -> str:137        out = self.llm.create_chat_completion(messages=messages, **kw)138        return out["choices"][0]["message"]["content"]139 140    @modal.method()141    def complete_json(self, messages: list[dict], schema: dict, **kw) -> dict:142        import json143 144        prompt_chars = sum(len(m.get("content", "")) for m in messages)145        print(f"[modal] LLM complete_json: {len(messages)} msgs, {prompt_chars} chars in")146        out = self.llm.create_chat_completion(147            messages=messages,148            response_format={"type": "json_object", "schema": schema},149            temperature=kw.get("temperature", 0.7),150            top_p=kw.get("top_p", 0.9),151            max_tokens=kw.get("max_tokens", 512),152            presence_penalty=kw.get("presence_penalty", 0.0),153        )154 155        content = out["choices"][0]["message"]["content"]156        print(f"[modal] LLM complete_json: {len(content)} chars out")157        return json.loads(content)158 159 160# --------------------------------------------------------------------------- #161#  Painter backend  (SDXL-base-1.0 + Lightning LoRA on A10G)162# --------------------------------------------------------------------------- #163 164 165@app.cls(166    image=painter_image,167    gpu="A10G",168    volumes={IMAGE_DIR: image_volume},169    timeout=120,170    scaledown_window=600,  # keep warm 10 min — avoids cold-start during a play session171)172class ModalPainterBackend:173    """SDXL-base-1.0 + ByteDance SDXL-Lightning 4-step LoRA.174 175    Mirrors the local SdxlLightningPainter exactly:176      - EulerDiscreteScheduler with trailing timestep spacing177      - LoRA fused into weights (fuse_lora) for faster inference178      - guidance_scale=0.0 (required by Lightning distillation)179      - 4 inference steps180      - Optional rembg background removal for sprites181    """182 183    @modal.enter()184    def load(self) -> None:185        import torch186        from diffusers import AutoencoderKL, EulerDiscreteScheduler, StableDiffusionXLPipeline187 188        print(f"[modal] Loading SDXL-base from {IMAGE_DIR}/{SDXL_BASE_LOCAL} ...")189        # fp16-safe VAE: decodes natively in fp16 (no "weird pixels", no per-render190        # fp32 upcast like the deprecated pipe.upcast_vae() workaround)191        vae = AutoencoderKL.from_pretrained(192            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16193        )194        self.pipe = StableDiffusionXLPipeline.from_pretrained(195            f"{IMAGE_DIR}/{SDXL_BASE_LOCAL}",196            torch_dtype=torch.float16,197            variant="fp16",198            vae=vae,199        ).to("cuda")200 201        # Lightning requires EulerDiscrete with trailing timestep spacing202        self.pipe.scheduler = EulerDiscreteScheduler.from_config(203            self.pipe.scheduler.config, timestep_spacing="trailing"204        )205 206        # Load Lightning LoRA from HF Hub (~400 MB, fast)207        print(f"[modal] Loading Lightning LoRA from {LIGHTNING_LORA_REPO} ...")208        self.pipe.load_lora_weights(LIGHTNING_LORA_REPO, weight_name=LIGHTNING_LORA_FILE)209        self.pipe.fuse_lora()  # bake into weights for faster inference210 211        self.torch = torch212        print("[modal] SDXL-Lightning ready on GPU")213 214    @modal.method()215    def render(216        self,217        prompt: str,218        negative_prompt: str,219        seed: int,220        size: int,221        steps: int,222        guidance_scale: float = 0.0,223        remove_bg: bool = False,224    ) -> bytes:225        """Generate an image and return PNG bytes.226 227        Args:228            prompt: Positive prompt.229            negative_prompt: Negative prompt.230            seed: RNG seed for reproducibility.231            size: Square image side in pixels.232            steps: Inference steps (4 for Lightning).233            guidance_scale: CFG scale — 0.0 for Lightning sprites, >1.0 for backdrops.234            remove_bg: Run rembg on the output (sprites only).235        """236        import io237 238        gen = self.torch.Generator(device="cuda").manual_seed(seed)239        result = self.pipe(240            prompt=prompt,241            negative_prompt=negative_prompt or None,242            num_inference_steps=steps,243            guidance_scale=guidance_scale,244            height=size,245            width=size,246            generator=gen,247        )248        img = result.images[0]249 250        if remove_bg:251            from rembg import new_session, remove  # noqa: PLC0415252 253            # Reuse one ONNX session per container — remove() without a session254            # reloads the ~170MB u2net model on every sprite. Lazy (not in255            # @modal.enter()) so backdrop-only requests never pay for it.256            if getattr(self, "_rembg_session", None) is None:257                self._rembg_session = new_session()258            img = remove(img, session=self._rembg_session)  # returns RGBA PIL image259 260        buf = io.BytesIO()261        img.save(buf, format="PNG")  # PNG supports RGBA transparency262        return buf.getvalue()263 264 265# =========================================================================== #266#  Smoke tests267# =========================================================================== #268 269 270@app.local_entrypoint()271def smoke() -> None:272    backend = ModalLLMBackend()273    reply = backend.complete.remote(274        [{"role": "user", "content": "Say hello in one word."}], max_tokens=10275    )276    print("LLM smoke:", reply)277 278 279@app.local_entrypoint()280def smoke_painter() -> None:281    import io282 283    from PIL import Image284 285    backend = ModalPainterBackend()286 287    # Backdrop test288    png = backend.render.remote(289        prompt="Japanese anime forest, glowing mushrooms, painterly background",290        negative_prompt="text, watermark, characters, person",291        seed=42,292        size=512,293        steps=4,294        guidance_scale=0.0,295        remove_bg=False,296    )297    img = Image.open(io.BytesIO(png))298    img.save("smoke_backdrop.png")299    print(f"Backdrop smoke OK -> smoke_backdrop.png  {img.size}")300 301    # Sprite test (with rembg background removal)302    png_sprite = backend.render.remote(303        prompt="anime girl, school uniform, happy expression, white background",304        negative_prompt="text, watermark, scenery, complex background",305        seed=7,306        size=512,307        steps=4,308        guidance_scale=0.0,309        remove_bg=True,310    )311    img_sprite = Image.open(io.BytesIO(png_sprite))312    img_sprite.save("smoke_sprite.png")313    print(f"Sprite smoke OK  -> smoke_sprite.png   {img_sprite.size}  mode={img_sprite.mode}")314