cpu_encoder.py

# ==========================================================================
# پروسه‌ی مستقل انکودرهای CPU (T5 متن، CLIP تصویر، پیش‌پردازش پیکسل‌ها)
# ==========================================================================
# چرا پروسه‌ی جدا؟
#   در ZeroGPU، پنجره‌ی GPU با fork از پروسه‌ی اصلی ساخته می‌شود. اجرای مدل‌های
#   transformers روی CPU داخل پروسه‌ی اصلی، تابع torch.cuda.is_current_stream_capturing()
#   را صدا می‌زند؛ همین یک فراخوانی وضعیت CUDA پروسه‌ی اصلی را خراب می‌کند و از آن به بعد
#   هر پنجره‌ی GPU با خطای «RuntimeError: No CUDA GPUs are available» شکست می‌خورد.
#   با اجرای همه‌ی محاسبات CPU در این پروسه‌ی مستقل (که spaces را import نمی‌کند و GPU
#   نمی‌بیند)، پروسه‌ی اصلی هیچ‌وقت آلوده نمی‌شود.
#
# پروتکل: هر درخواست یک خط JSON روی stdin، هر پاسخ یک خط JSON روی stdout.
# تنسورها از طریق فایل (torch.save) جابه‌جا می‌شوند.
# ==========================================================================
import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = ""

# stdout فقط برای پروتکل؛ هر چاپ دیگری (کتابخانه‌ها) به stderr (لاگ اسپیس) می‌رود
_PROTO = os.fdopen(os.dup(1), "w", buffering=1, encoding="utf-8")
os.dup2(2, 1)
sys.stdout = sys.stderr

import json
import traceback

import cv2
import torch
from PIL import Image
from transformers import AutoTokenizer, CLIPVisionModel, UMT5EncoderModel

from diffusers.modular_pipelines.wan_animate_2.encoders import clip_visual_encode
from diffusers.modular_pipelines.wan_animate_2.video_processor import WanAnimate2VideoProcessor

import animate2_engine as eng

# پرامپت پیش‌فرض رسمی Wan2.2-Animate-14B (حالت جایگزینی شخصیت)
MIX_DEFAULT_PROMPT = "视频中的人在做动作"


def log(msg):
    print(f"[cpu-encoder] {msg}", file=sys.stderr, flush=True)


def reply(obj):
    _PROTO.write(json.dumps(obj, ensure_ascii=False) + "\n")
    _PROTO.flush()


model_dir = sys.argv[1]
fixed_out = sys.argv[2]

tokenizer = AutoTokenizer.from_pretrained(model_dir, subfolder="tokenizer")
text_encoder = UMT5EncoderModel.from_pretrained(model_dir, subfolder="text_encoder", dtype=torch.bfloat16).eval()
image_encoder = CLIPVisionModel.from_pretrained(model_dir, subfolder="image_encoder", dtype=torch.float32).eval()

image_processor = WanAnimate2VideoProcessor(vae_scale_factor=8, spatial_patch_size=(2, 2), resample="bicubic")
video_processor = WanAnimate2VideoProcessor(vae_scale_factor=8, spatial_patch_size=(2, 2), resample="bilinear")


_mix_image_encoder = None


def get_mix_image_encoder(clip_dir):
    """انکودر تصویر مدل Wan2.2-Animate-14B؛ فقط در اولین درخواست حالت جایگزینی بارگذاری می‌شود."""
    global _mix_image_encoder
    if _mix_image_encoder is None:
        log("loading mix image encoder...")
        _mix_image_encoder = CLIPVisionModel.from_pretrained(
            clip_dir, subfolder="image_encoder", dtype=torch.float32
        ).eval()
    return _mix_image_encoder


def encode_text(prompt):
    global text_encoder
    try:
        return eng.encode_prompt_t5(text_encoder, tokenizer, prompt)
    except Exception as e:
        log(f"T5 با bf16 روی CPU ناموفق بود ({e})؛ سوییچ به float32")
        text_encoder = text_encoder.to(torch.float32)
        return eng.encode_prompt_t5(text_encoder, tokenizer, prompt)


def read_frames(video_path, indices):
    wanted = set(indices)
    last = max(wanted)
    cache = {}
    cap = cv2.VideoCapture(video_path)
    idx = 0
    while idx <= last:
        if idx in wanted:
            ok, frame = cap.read()
            if not ok:
                break
            cache[idx] = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        else:
            if not cap.grab():
                break
        idx += 1
    cap.release()
    if not cache:
        raise RuntimeError("خواندن فریم‌های ویدیو ناموفق بود.")
    fallback = cache[max(cache)]
    return [cache.get(i, fallback) for i in indices]


def atomic_save(obj, path):
    tmp = path + ".tmp"
    torch.save(obj, tmp)
    os.replace(tmp, path)


@torch.no_grad()
def cmd_prepare(req):
    image = Image.open(req["image"]).convert("RGB")
    img_h, img_w = image_processor.get_default_height_width(image)
    height, width, crop_region = eng.resolve_frame_size(img_h, img_w, int(req["area"]))

    image_pixels = image_processor.preprocess(image, height=height, width=width, resize_mode="fill").to(torch.float32)
    first_frame = read_frames(req["video"], [0])
    first_pixels = video_processor.preprocess_video(first_frame, height=height, width=width, resize_mode="fill")

    clip_ref = clip_visual_encode(image_encoder, image_pixels[0], "cpu", torch.float32)
    clip_drive = clip_visual_encode(image_encoder, first_pixels[0, :, 0].to(torch.float32), "cpu", torch.float32)

    prompt = (req.get("prompt") or "").strip()
    prompt_embeds = encode_text(prompt) if prompt else None

    atomic_save(
        {
            "image_pixels": image_pixels.to(torch.bfloat16),
            "clip_ref": clip_ref.to(torch.bfloat16),
            "clip_drive": clip_drive.to(torch.bfloat16),
            "prompt_embeds": prompt_embeds,
        },
        req["out"],
    )
    return {"height": int(height), "width": int(width), "crop_region": [int(v) for v in crop_region]}


@torch.no_grad()
def cmd_mix_prepare(req):
    """
    حالت جایگزینی: اندازه‌ی فریم (resize_by_area رسمی)، تصویر مرجع با padding_resize (INTER_AREA)،
    ویژگی CLIP تصویر مرجع، نسخه‌ی باکیفیت مرجع برای هم‌ترازی کادر، و (در صورت وجود) امبدینگ پرامپت.
    """
    import numpy as np
    import mix_engine as mix

    cap = cv2.VideoCapture(req["video"])
    ok, frame = cap.read()
    cap.release()
    if not ok:
        raise RuntimeError("خواندن فریم اول ویدیو ناموفق بود.")
    vid_h, vid_w = frame.shape[:2]
    height, width = mix.mix_frame_size(vid_w, vid_h, int(req["area"]))
    if height < 64 or width < 64:
        raise RuntimeError("ابعاد ویدیو نامعتبر است.")

    image = np.asarray(Image.open(req["image"]).convert("RGB"))
    ref_src = mix.reference_source(image)                           # نسخه‌ی باکیفیت برای هم‌ترازی روی GPU
    ref = mix.reference_padded(ref_src, height, width)              # uint8 H,W,3 (INTER_AREA)
    ref_t = torch.from_numpy(np.ascontiguousarray(ref)).permute(2, 0, 1).contiguous()  # 3,H,W
    ref_pm1 = ref_t.to(torch.float32) / 127.5 - 1.0
    encoder = get_mix_image_encoder(req["clip_dir"])
    clip_ref = clip_visual_encode(encoder, ref_pm1, "cpu", torch.float32)

    prompt = (req.get("prompt") or "").strip()
    prompt_embeds = encode_text(prompt) if prompt else None
    atomic_save(
        {
            "ref_pixels": ref_t,
            "clip_ref": clip_ref.to(torch.bfloat16),
            "ref_src": torch.from_numpy(np.ascontiguousarray(ref_src)),
            "prompt_embeds": prompt_embeds,
        },
        req["out"],
    )
    return {"height": int(height), "width": int(width)}


@torch.no_grad()
def cmd_segments(req):
    real = int(req["real_frame_len"])
    height, width = int(req["height"]), int(req["width"])
    all_indices = {}
    for k in req["segments"]:
        k = int(k)
        start = k * eng.EFFECTIVE_SEGMENT
        all_indices[k] = [eng.zigzag_index(p, real) for p in range(start, start + eng.SEGMENT_FRAME_LENGTH)]
    unique = sorted({i for v in all_indices.values() for i in v})
    frames = dict(zip(unique, read_frames(req["video"], unique)))
    paths = {}
    for k, idxs in all_indices.items():
        pixels = video_processor.preprocess_video(
            [frames[i] for i in idxs], height=height, width=width, resize_mode="fill"
        )
        path = os.path.join(req["workdir"], f"seg_{k}.pt")
        atomic_save(eng.pixels_to_uint8(pixels), path)
        del pixels
        paths[str(k)] = path
    return {"paths": paths}


def main():
    log("encoding fixed prompts...")
    atomic_save(
        {
            "default": encode_text(eng.DEFAULT_PROMPT),
            "negative": encode_text(eng.DEFAULT_NEGATIVE_PROMPT),
            "ref": encode_text(eng.DEFAULT_PROMPT_REF),
            "mix": encode_text(MIX_DEFAULT_PROMPT),
        },
        fixed_out,
    )
    reply({"ready": True})
    log("ready")

    handlers = {"prepare": cmd_prepare, "segments": cmd_segments, "mix_prepare": cmd_mix_prepare, "ping": lambda req: {}}
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            result = handlers[req["cmd"]](req)
            reply({"ok": True, **result})
        except Exception as e:
            traceback.print_exc()
            reply({"ok": False, "error": f"{type(e).__name__}: {e}"})


if __name__ == "__main__":
    main()