Restore original app; fix bucket path to /data/v1/model/...

app.py

@@ -1,495 +1,88 @@
-"""PinkCherry LTX 2.3 — Gradio Space for Quantumbraid/grok."""
-
-from __future__ import annotations
-
-import json
-import logging
-import os
-import random
-import struct
-import subprocess
-import sys
-import tempfile
-from pathlib import Path
-
-# Disable torch.compile before any torch import.
-os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
-os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
-
-# ---------------------------------------------------------------------------
-# Install LTX-2 packages (ltx-core + ltx-pipelines)
-# ---------------------------------------------------------------------------
-
-LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
-LTX_COMMIT_SHA = "780984275fd47128b02bef9b5c085404276866ee"  # main — has OffloadMode + blocks API
-LTX_REPO_DIR = Path(__file__).resolve().parent / "LTX-2"
-
-
-def _ensure_ltx_repo() -> None:
-    import shutil
-
-    if LTX_REPO_DIR.exists():
-        head = subprocess.run(
-            ["git", "-C", str(LTX_REPO_DIR), "rev-parse", "HEAD"],
-            capture_output=True,
-            text=True,
-            check=False,
-        )
-        if head.returncode == 0 and head.stdout.strip() == LTX_COMMIT_SHA:
-            return
-        shutil.rmtree(LTX_REPO_DIR, ignore_errors=True)
-
-    print(f"Cloning {LTX_REPO_URL} @ {LTX_COMMIT_SHA[:8]}...")
-    LTX_REPO_DIR.mkdir(parents=True, exist_ok=True)
-    subprocess.run(["git", "init", str(LTX_REPO_DIR)], check=True)
-    subprocess.run(["git", "-C", str(LTX_REPO_DIR), "remote", "add", "origin", LTX_REPO_URL], check=True)
-    subprocess.run(
-        ["git", "-C", str(LTX_REPO_DIR), "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA],
-        check=True,
-    )
-    subprocess.run(["git", "-C", str(LTX_REPO_DIR), "checkout", LTX_COMMIT_SHA], check=True)
-
-
-_ensure_ltx_repo()
-
-print("Installing ltx-core and ltx-pipelines...")
-subprocess.run(
-    [
-        sys.executable,
-        "-m",
-        "pip",
-        "install",
-        "--quiet",
-        "--no-deps",
-        "-e",
-        str(LTX_REPO_DIR / "packages" / "ltx-core"),
-        "-e",
-        str(LTX_REPO_DIR / "packages" / "ltx-pipelines"),
-    ],
-    check=True,
-)
-
-sys.path.insert(0, str(LTX_REPO_DIR / "packages" / "ltx-pipelines" / "src"))
-sys.path.insert(0, str(LTX_REPO_DIR / "packages" / "ltx-core" / "src"))
-
 import gradio as gr
-import numpy as np
-import spaces
+import os
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
-
-from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
-from ltx_core.loader.primitives import StateDict
-from ltx_core.loader.sft_loader import SafetensorsStateDictLoader
-from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
-from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
-from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
-from ltx_pipelines.utils.args import ImageConditioningInput
-from ltx_pipelines.utils.constants import (
-    DEFAULT_NEGATIVE_PROMPT,
-    LTX_2_3_PARAMS,
-    detect_params,
-)
-from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import OffloadMode
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# ---------------------------------------------------------------------------
-# Bucket paths (mounted at /data)
-# ---------------------------------------------------------------------------
-
-DATA_ROOT = Path(os.environ.get("LTX_DATA_ROOT", "/data"))
-CHECKPOINT_PATH = DATA_ROOT / "v1/model/SexGod_PinkCherry_dev_bf16_LTX23_v1.safetensors"
-DISTILLED_LORA_PATH = (
-    DATA_ROOT / "v1/distil_lora/ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors"
+from diffusers import LTXPipeline  # LTX video generation pipeline (correct class)
+
+# -------------------------------------------------
+# Load the model (from the attached bucket)
+# -------------------------------------------------
+
+# Mounted bucket path inside the Space container
+DATA_ROOT = os.environ.get("LTX_DATA_ROOT", "/data")
+model_path = os.path.join(
+    DATA_ROOT,
+    "v1",
+    "model",
+    "SexGod_PinkCherry_dev_bf16_LTX23_v1.safetensors",
 )
-WORKFLOW_PATH = DATA_ROOT / "v1/PInkCherry_LTX23_NSFW_Workflow.json"
-DISTILLED_LORA_STRENGTH = 0.6  # from bundled ComfyUI workflow
-
-CACHE_DIR = Path(os.environ.get("LTX_CACHE_DIR", str(Path.home() / ".cache" / "grok-ltx")))
-CACHE_DIR.mkdir(parents=True, exist_ok=True)
-
-SPATIAL_UPSAMPLER_PATH = CACHE_DIR / "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
-GEMMA_ROOT = CACHE_DIR / "gemma-3-12b-it"
-
-FRAME_RATE = 24.0
-MAX_SEED = np.iinfo(np.int32).max
-TOKEN_ENV_NAMES = ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGINGFACE_TOKEN")
-BUCKET_ID = "Quantumbraid/PinkCherry_NSFW_LTX23-bucket"
-
-RESOLUTION_PRESETS = {
-    "16:9 low (768×512)": (768, 512),
-    "16:9 medium (1024×576)": (1024, 576),
-    "16:9 high (1536×864)": (1536, 864),
-    "9:16 low (512×768)": (512, 768),
-    "9:16 medium (576×1024)": (576, 1024),
-    "9:16 high (864×1536)": (864, 1536),
-    "1:1 low (768×768)": (768, 768),
-}
-
-# ---------------------------------------------------------------------------
-# FUSE-safe safetensors loader (bucket mounts can deadlock on mmap)
-# ---------------------------------------------------------------------------
-
-_SAFETENSORS_DTYPE_MAP = {
-    "F64": torch.float64,
-    "F32": torch.float32,
-    "F16": torch.float16,
-    "BF16": torch.bfloat16,
-    "F8_E5M2": torch.float8_e5m2,
-    "F8_E4M3": torch.float8_e4m3fn,
-    "I64": torch.int64,
-    "I32": torch.int32,
-    "I16": torch.int16,
-    "I8": torch.int8,
-    "U8": torch.uint8,
-    "BOOL": torch.bool,
-}
-
-
-def _patched_safetensors_load(self, path, sd_ops, device=None):
-    sd = {}
-    size = 0
-    dtype = set()
-    device = device or torch.device("cpu")
-    model_paths = path if isinstance(path, list) else [path]
-    for shard_path in model_paths:
-        with open(shard_path, "rb") as f:
-            header_len = struct.unpack("<Q", f.read(8))[0]
-            header = json.loads(f.read(header_len).decode("utf-8"))
-            data_base = 8 + header_len
-            for name, meta in header.items():
-                if name == "__metadata__":
-                    continue
-                expected_name = name if sd_ops is None else sd_ops.apply_to_key(name)
-                if expected_name is None:
-                    continue
-                start, end = meta["data_offsets"]
-                f.seek(data_base + start)
-                buf = f.read(end - start)
-                t = torch.frombuffer(
-                    bytearray(buf), dtype=_SAFETENSORS_DTYPE_MAP[meta["dtype"]]
-                ).reshape(meta["shape"])
-                t = t.to(device=device, non_blocking=True, copy=False)
-                kvs = (
-                    ((expected_name, t),)
-                    if sd_ops is None
-                    else sd_ops.apply_to_key_value(expected_name, t)
-                )
-                for key, v in kvs:
-                    size += v.nbytes
-                    dtype.add(v.dtype)
-                    sd[key] = v
-    return StateDict(sd=sd, device=device, size=size, dtype=dtype)
-
-
-SafetensorsStateDictLoader.load = _patched_safetensors_load
-
-
-# ---------------------------------------------------------------------------
-# xformers attention patch
-# ---------------------------------------------------------------------------
-
-from ltx_core.model.transformer import attention as _attn_mod
-
-try:
-    from xformers.ops import memory_efficient_attention as _mea
-
-    _attn_mod.memory_efficient_attention = _mea
-    try:
-        from xformers.ops.fmha import _set_use_fa3
-
-        _set_use_fa3(False)
-    except Exception:
-        pass
-except Exception as exc:
-    logger.warning("xformers patch skipped: %s", exc)
-
-torch._dynamo.config.suppress_errors = True
-torch._dynamo.config.disable = True
-
-
-# ---------------------------------------------------------------------------
-# Asset verification + hub downloads
-# ---------------------------------------------------------------------------
-
-
-def _get_hf_token() -> str | None:
-    for name in TOKEN_ENV_NAMES:
-        token = os.environ.get(name)
-        if token and token.strip():
-            return token.strip()
-    return None
-
 
-def _token_status() -> str:
-    present = [name for name in TOKEN_ENV_NAMES if os.environ.get(name)]
-    return present[0] if present else "none"
-
-
-def _describe_data_root() -> str:
-    if not DATA_ROOT.exists():
-        return f"{DATA_ROOT} does not exist (bucket volume not mounted?)"
-    try:
-        entries = sorted(p.name for p in DATA_ROOT.iterdir())
-    except OSError as exc:
-        return f"{DATA_ROOT} exists but is not readable: {exc}"
-    if not entries:
-        return f"{DATA_ROOT} is empty"
-    preview = ", ".join(entries[:12])
-    suffix = " ..." if len(entries) > 12 else ""
-    return f"{DATA_ROOT} contains: {preview}{suffix}"
-
-
-def _require_bucket_files() -> None:
-    missing = [p for p in (CHECKPOINT_PATH, DISTILLED_LORA_PATH) if not p.exists()]
-    if missing:
-        token_hint = (
-            f"A Space secret named HF_TOKEN is set ({_token_status()})."
-            if _get_hf_token()
-            else (
-                "No HF token secret found. Add a Space secret named HF_TOKEN with a read token "
-                f"from the account that owns {BUCKET_ID}, then restart the Space."
-            )
-        )
-        raise FileNotFoundError(
-            "Missing bucket files. Mount the bucket at /data and verify the v1/ layout:\n"
-            + "\n".join(f"  - {p}" for p in missing)
-            + f"\n\nBucket: {BUCKET_ID}\n"
-            + f"Mount: hf spaces volumes set Quantumbraid/tst "
-            + f"-v hf://buckets/{BUCKET_ID}:/data\n"
-            + f"Diagnostics: {_describe_data_root()}\n"
-            + token_hint
-        )
-
-
-def _ensure_hub_assets() -> None:
-    token = _get_hf_token()
-    if token is None:
-        logger.warning(
-            "No HF token secret detected (%s). Gated Hub downloads may fail.",
-            _token_status(),
-        )
-
-    if not SPATIAL_UPSAMPLER_PATH.exists():
-        logger.info("Downloading spatial upsampler...")
-        hf_hub_download(
-            repo_id="Lightricks/LTX-2.3",
-            filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
-            token=token,
-            local_dir=str(CACHE_DIR),
-        )
-
-    if not GEMMA_ROOT.exists() or not any(GEMMA_ROOT.rglob("model*.safetensors")):
-        logger.info("Downloading Gemma 3 text encoder...")
-        snapshot_download(
-            "google/gemma-3-12b-it-qat-q4_0-unquantized",
-            token=token,
-            local_dir=str(GEMMA_ROOT),
-        )
-
-
-_require_bucket_files()
-_ensure_hub_assets()
-
-PIPELINE_PARAMS = detect_params(str(CHECKPOINT_PATH))
-
-distilled_lora = [
-    LoraPathStrengthAndSDOps(
-        str(DISTILLED_LORA_PATH),
-        DISTILLED_LORA_STRENGTH,
-        LTXV_LORA_COMFY_RENAMING_MAP,
+if not os.path.isfile(model_path):
+    raise FileNotFoundError(
+        f"Checkpoint not found at {model_path}. "
+        f"Ensure bucket Quantumbraid/PinkCherry_NSFW_LTX23-bucket is mounted at {DATA_ROOT}."
     )
-]
 
-logger.info("Loading TI2VidTwoStagesPipeline (FP8 + CPU offload)...")
-pipeline = TI2VidTwoStagesPipeline(
-    checkpoint_path=str(CHECKPOINT_PATH),
-    distilled_lora=distilled_lora,
-    spatial_upsampler_path=str(SPATIAL_UPSAMPLER_PATH),
-    gemma_root=str(GEMMA_ROOT),
-    loras=[],
-    quantization=build_fp8_cast_policy(str(CHECKPOINT_PATH)),
-    offload_mode=OffloadMode.CPU,
+# Initialise the LTX pipeline from the local checkpoint
+pipe = LTXPipeline.from_pretrained(
+    model_path,
+    torch_dtype=torch.float16,
 )
-logger.info("Pipeline ready.")
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _snap_frames(duration: float, fps: float = FRAME_RATE) -> int:
-    raw = int(duration * fps) + 1
-    k = max(0, (raw - 1 + 7) // 8)
-    return 8 * k + 1
-
-
-def _resolution_from_preset(preset: str) -> tuple[int, int]:
-    return RESOLUTION_PRESETS.get(preset, (768, 512))
-
-
-@spaces.GPU(duration=120)
-@torch.inference_mode()
-def generate(
-    prompt: str,
-    negative_prompt: str,
-    input_image,
-    duration: float,
-    resolution_preset: str,
-    num_inference_steps: int,
-    enhance_prompt: bool,
-    seed: int,
-    randomize_seed: bool,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-
-    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    width, height = _resolution_from_preset(resolution_preset)
-    num_frames = _snap_frames(duration)
-
-    images: list[ImageConditioningInput] = []
-    if input_image is not None:
-        temp_dir = Path("inputs")
-        temp_dir.mkdir(exist_ok=True)
-        image_path = temp_dir / f"frame_{current_seed}.jpg"
-        if hasattr(input_image, "save"):
-            input_image.save(image_path)
-        else:
-            image_path = Path(input_image)
-        images = [ImageConditioningInput(path=str(image_path), frame_idx=0, strength=1.0)]
-
-    logger.info(
-        "Generating %dx%d, %d frames (%.1fs), seed=%d, steps=%d",
-        width,
-        height,
-        num_frames,
-        duration,
-        current_seed,
-        num_inference_steps,
-    )
-
-    tiling_config = TilingConfig.default()
-    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
-
-    video, audio = pipeline(
+pipe = pipe.to("cuda")
+
+# -------------------------------------------------
+# Gradio UI helper
+# -------------------------------------------------
+def generate_video(prompt: str, num_frames: int = 16, seed: int = 42):
+    """Generate a short video (as a GIF) from a text prompt."""
+    generator = torch.Generator("cuda").manual_seed(seed)
+
+    # Run the pipeline – the return type can differ across versions,
+    # so we handle both the new attribute style and the older tuple style.
+    output = pipe(
         prompt=prompt,
-        negative_prompt=negative_prompt or DEFAULT_NEGATIVE_PROMPT,
-        seed=current_seed,
-        height=height,
-        width=width,
         num_frames=num_frames,
-        frame_rate=FRAME_RATE,
-        num_inference_steps=num_inference_steps,
-        video_guider_params=PIPELINE_PARAMS.video_guider_params,
-        audio_guider_params=PIPELINE_PARAMS.audio_guider_params,
-        images=images,
-        tiling_config=tiling_config,
-        enhance_prompt=enhance_prompt,
+        generator=generator,
     )
-
-    output_path = tempfile.mktemp(suffix=".mp4")
-    encode_video(
-        video=video,
-        fps=FRAME_RATE,
-        audio=audio,
-        output_path=output_path,
-        video_chunks_number=video_chunks_number,
-    )
-    return output_path, current_seed
-
-
-# ---------------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------------
-
-DEFAULT_PROMPT = (
-    "Cinematic close-up, natural warm lighting, smooth camera motion, "
-    "detailed textures, synchronized ambient audio."
+    # Newer versions return an object with a .videos attribute.
+    if hasattr(output, "videos"):
+        video = output.videos
+    else:
+        # Older versions return a tuple where the first element is the video tensor.
+        video = output[0]
+
+    # video shape: (B, T, C, H, W). We take the first batch element.
+    frames_tensor = video[0]                     # (T, C, H, W)
+    frames = [
+        frame.cpu().numpy().transpose(1, 2, 0)   # convert to H×W×C for Gradio/GIF
+        for frame in frames_tensor
+    ]
+    return frames
+
+# -------------------------------------------------
+# Build the Gradio interface
+# -------------------------------------------------
+title = "💥 PinkCherry NSFW LTX 2.3 – Video Generation Demo"
+description = (
+    "Enter a prompt (NSFW/uncensored) and generate a short video. "
+    "⚠️ **Not safe for work – use responsibly**."
 )
 
-with gr.Blocks(title="Grok LTX 2.3") as demo:
-    gr.Markdown(
-        "# Grok — PinkCherry LTX 2.3\n"
-        "Text-to-video and image-to-video with synchronized audio. "
-        "Checkpoint and distilled LoRA load from the mounted `/data` bucket."
-    )
-
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="First frame (optional, image-to-video)", type="pil")
-            prompt = gr.Textbox(
-                label="Prompt",
-                lines=4,
-                value=DEFAULT_PROMPT,
-                placeholder="Describe scene, motion, camera, lighting, and audio...",
-            )
-            negative_prompt = gr.Textbox(
-                label="Negative prompt",
-                lines=2,
-                value=DEFAULT_NEGATIVE_PROMPT,
-            )
-
-            with gr.Row():
-                duration = gr.Slider(
-                    label="Duration (seconds)",
-                    minimum=2.0,
-                    maximum=10.0,
-                    value=5.0,
-                    step=0.5,
-                )
-                resolution_preset = gr.Dropdown(
-                    label="Resolution",
-                    choices=list(RESOLUTION_PRESETS.keys()),
-                    value="16:9 low (768×512)",
-                )
-
-            with gr.Row():
-                num_inference_steps = gr.Slider(
-                    label="Inference steps (stage 1)",
-                    minimum=8,
-                    maximum=40,
-                    value=PIPELINE_PARAMS.num_inference_steps,
-                    step=1,
-                )
-                enhance_prompt = gr.Checkbox(label="Enhance prompt", value=False)
-
-            with gr.Accordion("Advanced", open=False):
-                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            generate_btn = gr.Button("Generate video", variant="primary")
-
-        with gr.Column():
-            output_video = gr.Video(label="Output", autoplay=True)
-
-    gr.Markdown(
-        f"**Bucket checkpoint:** `{CHECKPOINT_PATH.name}`  \n"
-        f"**Distilled LoRA:** `{DISTILLED_LORA_PATH.name}` @ strength {DISTILLED_LORA_STRENGTH}  \n"
-        f"**Workflow reference:** `{WORKFLOW_PATH.name}`"
-    )
-
-    generate_btn.click(
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            input_image,
-            duration,
-            resolution_preset,
-            num_inference_steps,
-            enhance_prompt,
-            seed,
-            randomize_seed,
-        ],
-        outputs=[output_video, seed],
-    )
+demo = gr.Interface(
+    fn=generate_video,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="e.g. a futuristic city at night"),
+        gr.Slider(8, 32, step=4, label="Number of frames", value=16),
+        gr.Number(label="Seed (optional)", value=42),
+    ],
+    outputs=gr.Gallery(label="Generated video (GIF)"),
+    title=title,
+    description=description,
+    allow_flagging="never",
+    examples=[
+        ["a seductive dancer in a neon club", 16, 123],
+        ["a cyber-punk couple on a rooftop", 24, 456],
+    ],
+)
 
 if __name__ == "__main__":
     demo.launch()