GhostPackDemo/ghostpack.py

# ==========================================================
# 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()