# ========================================================================== # پروسه‌ی مستقل انکودرهای 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()