Create ghostpack.py

GhostPackDemo/ghostpack.py

@@ -0,0 +1,232 @@
+# ==========================================================
+# FILE: ghostpack.py
+# ==========================================================
+#!/usr/bin/env python3
+# ---------------------------------------------------------------------------
+# RELEASE – GhostPack Image-to-Video Generator
+# ---------------------------------------------------------------------------
+import os, sys, argparse, traceback
+import numpy as np, torch, einops, gradio as gr
+from PIL import Image
+from diffusers_helper.hf_login import login
+from diffusers import AutoencoderKLHunyuanVideo
+from transformers import (
+    LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer,
+    SiglipImageProcessor, SiglipVisionModel,
+)
+from diffusers_helper.hunyuan import (
+    encode_prompt_conds, vae_encode, vae_decode, vae_decode_fake,
+)
+from diffusers_helper.utils import (
+    save_bcthw_as_mp4, crop_or_pad_yield_mask, soft_append_bcthw,
+    resize_and_center_crop, generate_timestamp,
+)
+from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
+from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
+from diffusers_helper.memory import (
+    gpu, get_cuda_free_memory_gb, DynamicSwapInstaller,
+    unload_complete_models, load_model_as_complete,
+    fake_diffusers_current_device, move_model_to_device_with_memory_preservation,
+    offload_model_from_device_for_memory_preservation,
+)
+from diffusers_helper.thread_utils import AsyncStream, async_run
+from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
+from diffusers_helper.clip_vision import hf_clip_vision_encode
+from diffusers_helper.bucket_tools import find_nearest_bucket
+
+BASE  = os.path.abspath(os.path.dirname(__file__))
+CACHE = os.path.join(BASE, "hf_download")
+os.makedirs(CACHE, exist_ok=True)
+for v in ("HF_HOME", "TRANSFORMERS_CACHE", "HF_DATASETS_CACHE"): os.environ[v] = CACHE
+os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
+
+p = argparse.ArgumentParser()
+p.add_argument("--share", action="store_true")
+p.add_argument("--server", default="0.0.0.0")
+p.add_argument("--port", type=int, default=7860)
+p.add_argument("--inbrowser", action="store_true")
+args = p.parse_args()
+
+free_gb = get_cuda_free_memory_gb(gpu)
+hi_vram = free_gb > 60
+print(f"[GhostPack] Free VRAM: {free_gb:.1f} GB | High-VRAM: {hi_vram}")
+
+def llm(sub):  return LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder=sub, torch_dtype=torch.float16).cpu().eval()
+def clip(sub): return CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder=sub, torch_dtype=torch.float16).cpu().eval()
+
+text_enc  = llm("text_encoder")
+text_enc2 = clip("text_encoder_2")
+tok       = LlamaTokenizerFast.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder="tokenizer")
+tok2      = CLIPTokenizer.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder="tokenizer_2")
+vae       = AutoencoderKLHunyuanVideo.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder="vae", torch_dtype=torch.float16).cpu().eval()
+feat_ext  = SiglipImageProcessor.from_pretrained("lllyasviel/flux_redux_bfl", subfolder="feature_extractor")
+img_enc   = SiglipVisionModel.from_pretrained("lllyasviel/flux_redux_bfl", subfolder="image_encoder", torch_dtype=torch.float16).cpu().eval()
+trans     = HunyuanVideoTransformer3DModelPacked.from_pretrained("lllyasviel/FramePackI2V_HY", torch_dtype=torch.bfloat16).cpu().eval()
+trans.high_quality_fp32_output_for_inference = True
+
+if not hi_vram:
+    vae.enable_slicing(); vae.enable_tiling()
+else:
+    for m in (text_enc, text_enc2, img_enc, vae, trans): m.to(gpu)
+
+trans.to(dtype=torch.bfloat16)
+for m in (vae, img_enc, text_enc, text_enc2): m.to(dtype=torch.float16)
+for m in (vae, img_enc, text_enc, text_enc2, trans): m.requires_grad_(False)
+
+if not hi_vram:
+    DynamicSwapInstaller.install_model(trans, device=gpu)
+    DynamicSwapInstaller.install_model(text_enc, device=gpu)
+
+OUT = os.path.join(BASE, "outputs")
+os.makedirs(OUT, exist_ok=True)
+stream = AsyncStream()
+
+@torch.no_grad()
+def worker(img, p, n_p, sd, secs, win, stp, cfg, gsc, rsc, keep, tea, crf):
+    sections = max(round((secs*30)/(win*4)), 1)
+    job = generate_timestamp()
+    stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"Start"))))
+    try:
+        if not hi_vram: unload_complete_models(text_enc, text_enc2, img_enc, vae, trans)
+        stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"Text enc"))))
+        if not hi_vram:
+            fake_diffusers_current_device(text_enc, gpu)
+            load_model_as_complete(text_enc2, gpu)
+        lv, cp = encode_prompt_conds(p, text_enc, text_enc2, tok, tok2)
+        lv_n, cp_n = (torch.zeros_like(lv), torch.zeros_like(cp)) if cfg==1 else encode_prompt_conds(n_p, text_enc, text_enc2, tok, tok2)
+        lv, m    = crop_or_pad_yield_mask(lv,512)
+        lv_n, m_n= crop_or_pad_yield_mask(lv_n,512)
+        stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"Image"))))
+        H,W,_ = img.shape; h,w = find_nearest_bucket(H,W,640)
+        img_np = resize_and_center_crop(img,w,h)
+        Image.fromarray(img_np).save(os.path.join(OUT,f"{job}.png"))
+        img_pt = torch.from_numpy(img_np).float()/127.5-1; img_pt = img_pt.permute(2,0,1)[None,:,None]
+        stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"VAE"))))
+        if not hi_vram: load_model_as_complete(vae, gpu)
+        start_lat = vae_encode(img_pt, vae)
+        stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"Vision"))))
+        if not hi_vram: load_model_as_complete(img_enc, gpu)
+        img_hidden = hf_clip_vision_encode(img_np, feat_ext, img_enc).last_hidden_state
+        to = trans.dtype
+        lv, lv_n, cp, cp_n, img_hidden = (x.to(to) for x in (lv, lv_n, cp, cp_n, img_hidden))
+        stream.output_queue.push(("progress",(None,"",make_progress_bar_html(0,"Sample"))))
+        gen = torch.Generator("cpu").manual_seed(sd)
+        frames = win*4-3
+        hist_lat = torch.zeros((1,16,1+2+16,h//8,w//8), dtype=torch.float32).cpu()
+        hist_px=None; total=0
+        pad_seq=[3]+[2]*(sections-3)+[1,0] if sections>4 else list(reversed(range(sections)))
+        for pad in pad_seq:
+            last = pad==0
+            if stream.input_queue.top()=="end": stream.output_queue.push(("end",None)); return
+            pad_sz=pad*win
+            idx=torch.arange(0,sum([1,pad_sz,win,1,2,16])).unsqueeze(0)
+            a,b,c,d,e,f = idx.split([1,pad_sz,win,1,2,16],1)
+            clean_idx = torch.cat([a,d],1)
+            pre=start_lat.to(hist_lat); post,two,four=hist_lat[:,:,:1+2+16].split([1,2,16],2)
+            clean=torch.cat([pre,post],2)
+            if not hi_vram:
+                unload_complete_models()
+                move_model_to_device_with_memory_preservation(trans,gpu,keep)
+            trans.initialize_teacache(tea,stp)
+            def cb(d):
+                pv = vae_decode_fake(d["denoised"])
+                pv = (pv*255).cpu().numpy().clip(0,255).astype(np.uint8)
+                pv = einops.rearrange(pv,"b c t h w->(b h)(t w)c")
+                cur = d["i"]+1
+                stream.output_queue.push(("progress",(pv,f"{total*4-3}f",make_progress_bar_html(int(100*cur/stp),f"{cur}/{stp}"))))
+                if stream.input_queue.top()=="end":
+                    stream.output_queue.push(("end",None)); raise KeyboardInterrupt
+            new_lat = sample_hunyuan(
+                transformer=trans,sampler="unipc",width=w,height=h,frames=frames,
+                real_guidance_scale=cfg,distilled_guidance_scale=gsc,guidance_rescale=rsc,
+                num_inference_steps=stp,generator=gen,
+                prompt_embeds=lv,prompt_embeds_mask=m,prompt_poolers=cp,
+                negative_prompt_embeds=lv_n,negative_prompt_embeds_mask=m_n,negative_prompt_poolers=cp_n,
+                device=gpu,dtype=torch.bfloat16,image_embeddings=img_hidden,
+                latent_indices=c,clean_latents=clean,clean_latent_indices=clean_idx,
+                clean_latents_2x=two,clean_latent_2x_indices=e,clean_latents_4x=four,clean_latent_4x_indices=f,
+                callback=cb,
+            )
+            if last: new_lat=torch.cat([start_lat.to(new_lat),new_lat],2)
+            total+=new_lat.shape[2]; hist_lat=torch.cat([new_lat.to(hist_lat),hist_lat],2)
+            if not hi_vram:
+                offload_model_from_device_for_memory_preservation(trans,gpu,8)
+                load_model_as_complete(vae,gpu)
+            real=hist_lat[:,:,:total]
+            if hist_px is None:
+                hist_px = vae_decode(real,vae).cpu()
+            else:
+                sec_lat=win*2+1 if last else win*2
+                cur_px = vae_decode(real[:,:,:sec_lat],vae).cpu()
+                hist_px = soft_append_bcthw(cur_px,hist_px,win*4-3)
+            if not hi_vram: unload_complete_models()
+            mp4=os.path.join(OUT,f"{job}_{total}.mp4")
+            save_bcthw_as_mp4(hist_px,mp4,fps=30,crf=crf)
+            stream.output_queue.push(("file",mp4))
+            if last: break
+    except Exception:
+        traceback.print_exc(); stream.output_queue.push(("end",None))
+
+def ui():
+    css = make_progress_bar_css()+"""
+    body,.gradio-container,.gr-block{background:#121212;color:#eee}
+    .gr-button,.gr-button-primary{background:#006400;border:#006400}
+    .gr-button:hover,.gr-button-primary:hover{background:#00aa00;border:#00aa00}
+    input,textarea,.gr-input,.gr-textbox,.gr-slider,.gr-number{background:#1e1e1e;color:#eee;border-color:#006400}
+    """
+    quick=[["The girl dances gracefully, with clear movements, full of charm."],
+           ["A character doing some simple body movements."]]
+    blk=gr.Blocks(css=css).queue()
+    with blk:
+        gr.Markdown("# 👻 GhostPack Demo")
+        with gr.Row():
+            with gr.Column():
+                img=gr.Image(sources="upload",type="numpy",label="Image",height=320)
+                prm=gr.Textbox(label="Prompt")
+                ds=gr.Dataset(samples=quick,label="Quick List",components=[prm])
+                ds.click(lambda x:x[0],inputs=[ds],outputs=prm)
+                with gr.Row():
+                    b_go=gr.Button("Start"); b_end=gr.Button("End",interactive=False)
+                with gr.Group():
+                    tea=gr.Checkbox(label="Use TeaCache",value=True)
+                    npr=gr.Textbox(label="Negative Prompt",visible=False)
+                    se=gr.Number(label="Seed",value=31337,precision=0)
+                    sec=gr.Slider(label="Video Length (s)",minimum=1,maximum=120,value=5,step=0.1)
+                    win=gr.Slider(label="Latent Window",minimum=1,maximum=33,value=9,step=1,visible=False)
+                    stp=gr.Slider(label="Steps",minimum=1,maximum=100,value=25,step=1)
+                    cfg=gr.Slider(label="CFG",minimum=1,maximum=32,value=1,step=0.01,visible=False)
+                    gsc=gr.Slider(label="Distilled CFG",minimum=1,maximum=32,value=10,step=0.01)
+                    rsc=gr.Slider(label="CFG Re-Scale",minimum=0,maximum=1,value=0,step=0.01,visible=False)
+                    kee=gr.Slider(label="GPU Keep (GB)",minimum=6,maximum=128,value=6,step=0.1)
+                    crf=gr.Slider(label="MP4 CRF",minimum=0,maximum=100,value=16,step=1)
+            with gr.Column():
+                pv=gr.Image(label="Next Latents",height=200,visible=False,interactive=False)
+                vid=gr.Video(label="Finished",autoplay=True,height=512,loop=True,show_share_button=False)
+                gr.Markdown("Ending actions appear first; wait for start.")
+                dsc=gr.Markdown("")
+                bar=gr.HTML("")
+                log=gr.Markdown("")
+        inputs=[img,prm,npr,se,sec,win,stp,cfg,gsc,rsc,kee,tea,crf]
+        def launch(*xs):
+            global stream
+            if xs[0] is None: raise gr.Error("Upload an image.")
+            yield None,None,"","","",gr.update(interactive=False),gr.update(interactive=True)
+            stream=AsyncStream()
+            async_run(worker,*xs)
+            out=None; log=""
+            while True:
+                flag,data=stream.output_queue.next()
+                if flag=="file":
+                    out=data
+                    yield out,gr.update(),gr.update(),gr.update(),log,gr.update(interactive=False),gr.update(interactive=True)
+                if flag=="progress":
+                    pv,desc,html=data; log=desc
+                    yield gr.update(),gr.update(visible=True,value=pv),desc,html,log,gr.update(interactive=False),gr.update(interactive=True)
+                if flag=="end":
+                    yield out,gr.update(visible=False),gr.update(),"",log,gr.update(interactive=True),gr.update(interactive=False); break
+        b_go.click(launch,inputs,[vid,pv,dsc,bar,log,b_go,b_end])
+        b_end.click(lambda: stream.input_queue.push("end"))
+    blk.launch(server_name=args.server,server_port=args.port,share=args.share,inbrowser=args.inbrowser)
+
+if __name__ == "__main__":
+    ui()