Create app.py

app.py

@@ -0,0 +1,226 @@
+import os
+import gc
+import torch
+import numpy as np
+from PIL import Image
+import imageio
+import gradio as gr
+from huggingface_hub import hf_hub_download
+
+# ComfyUI imports (assumes ComfyUI folder is dedicated in repo)
+from comfy import model_management  # may be needed for plugin system
+from nodes import (
+    CheckpointLoaderSimple,
+    CLIPLoader,
+    CLIPTextEncode,
+    VAELoader,
+    VAEDecode,
+    KSampler,
+)
+from custom_nodes.ComfyUI_GGUF.nodes import UnetLoaderGGUF
+from comfy_extras.nodes_hunyuan import EmptyHunyuanLatentVideo
+from comfy_extras.nodes_images import SaveAnimatedWEBP
+from comfy_extras.nodes_video import SaveWEBM
+
+# Globals
+unet_loader = None
+clip_loader = None
+clip_encode_positive = None
+clip_encode_negative = None
+vae_loader = None
+empty_latent_video = None
+ksampler = None
+vae_decode = None
+
+
+# ✅ Ensure models are available via HF hub or local
+def ensure_model(repo_id, filename, folder):
+    os.makedirs(f"ComfyUI/models/{folder}", exist_ok=True)
+    local_path = os.path.join("ComfyUI", "models", folder, filename)
+    if not os.path.isfile(local_path):
+        hf_hub_download(repo_id=repo_id, filename=filename, local_dir=os.path.dirname(local_path))
+    return local_path
+
+
+# 1️⃣ Initialize imports and model loader utilities
+def imports_initialization():
+    global unet_loader, clip_loader, clip_encode_positive, clip_encode_negative
+    global vae_loader, empty_latent_video, ksampler, vae_decode
+
+    unet_loader = UnetLoaderGGUF()
+    clip_loader = CLIPLoader()
+    clip_encode_positive = CLIPTextEncode()
+    clip_encode_negative = CLIPTextEncode()
+    vae_loader = VAELoader()
+    empty_latent_video = EmptyHunyuanLatentVideo()
+    ksampler = KSampler()
+    vae_decode = VAEDecode()
+
+    return "✅ Imports done and models initialized."
+
+
+# Clean GPU memory
+def clear_memory():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    for obj in list(globals().values()):
+        try:
+            if torch.is_tensor(obj) or (hasattr(obj, "data") and torch.is_tensor(obj.data)):
+                del obj
+        except:
+            pass
+    gc.collect()
+
+
+# Save utility functions
+def save_as_mp4(images, prefix, fps):
+    os.makedirs("output", exist_ok=True)
+    path = f"output/{prefix}.mp4"
+    writer = imageio.get_writer(path, fps=fps)
+    for img in images:
+        writer.append_data((img.cpu().numpy() * 255).astype(np.uint8))
+    writer.close()
+    return path
+
+def save_as_webm(images, prefix, fps):
+    os.makedirs("output", exist_ok=True)
+    path = f"output/{prefix}.webm"
+    writer = imageio.get_writer(
+        path, format='FFMPEG', fps=fps,
+        codec='vp9', quality=20
+    )
+    for img in images:
+        writer.append_data((img.cpu().numpy() * 255).astype(np.uint8))
+    writer.close()
+    return path
+
+def save_as_image(img, prefix):
+    os.makedirs("output", exist_ok=True)
+    path = f"output/{prefix}.png"
+    pil = Image.fromarray((img.cpu().numpy() * 255).astype(np.uint8))
+    pil.save(path)
+    return path
+
+
+# 2️⃣ Text-to-Video generation pipeline
+def generate_video(
+    positive_prompt, negative_prompt,
+    width, height, seed, steps, cfg_scale,
+    sampler_name, scheduler, frames, fps, output_format, use_q6
+):
+    log = []
+
+    # 2a. Download or load model files
+    unet_file = ensure_model(
+        "city96/Wan2.1-T2V-14B-gguf",
+        "wan2.1-t2v-14b-Q6_K.gguf" if use_q6 else "wan2.1-t2v-14b-Q5_0.gguf",
+        "unet"
+    )
+    text_enc_file = ensure_model(
+        "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+        "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
+        "text_encoders"
+    )
+    vae_file = ensure_model(
+        "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
+        "wan_2.1_vae.safetensors",
+        "vae"
+    )
+
+    # 2b. Encode text prompts
+    log.append("🔧 Encoding prompts...")
+    clip_model = clip_loader.load_clip(text_enc_file, "wan", "default")[0]
+    pos = clip_encode_positive.encode(clip_model, positive_prompt)[0]
+    neg = clip_encode_negative.encode(clip_model, negative_prompt)[0]
+    del clip_model
+    clear_memory()
+
+    # 2c. Setup latent video
+    latent = empty_latent_video.generate(width, height, frames, 1)[0]
+
+    # 2d. Sample using UNet
+    model = unet_loader.load_unet(unet_file)[0]
+    log.append("🎥 Sampling latents...")
+    sampled = ksampler.sample(
+        model=model,
+        seed=seed,
+        steps=steps,
+        cfg=cfg_scale,
+        sampler_name=sampler_name,
+        scheduler=scheduler,
+        positive=pos,
+        negative=neg,
+        latent_image=latent
+    )[0]
+    del model
+    clear_memory()
+
+    # 2e. Decode via VAE
+    log.append("🔓 Decoding with VAE...")
+    vae_model = vae_loader.load_vae(vae_file)[0]
+    decoded = vae_decode.decode(vae_model, sampled)[0]
+    del vae_model
+    clear_memory()
+
+    # 2f. Save output
+    filename = "hf_gen"
+    if frames == 1:
+        log.append("💾 Saving single frame...")
+        out = save_as_image(decoded[0], filename)
+    else:
+        if output_format == "webm":
+            log.append("💾 Saving as WEBM...")
+            out = save_as_webm(decoded, filename, fps)
+        else:
+            log.append("💾 Saving as MP4...")
+            out = save_as_mp4(decoded, filename, fps)
+
+    log.append(f"✅ Saved: {out}")
+    clear_memory()
+    return "\n".join(log), out
+
+
+# 3️⃣ Gradio UI
+
+app = gr.Blocks()
+with app:
+    gr.Markdown("# ComfyUI Text‑to‑Video on Hugging Face Spaces")
+
+    with gr.Tab("Initialize"):
+        init_btn = gr.Button("Initialize Models")
+        init_out = gr.Textbox(lines=3, interactive=False, label="Status")
+        init_btn.click(imports_initialization, None, init_out)
+
+    with gr.Tab("Generate"):
+        with gr.Row():
+            pos = gr.Textbox(label="Positive Prompt", value="lion")
+            neg = gr.Textbox(label="Negative Prompt", value="")
+        with gr.Row():
+            w = gr.Slider(64, 1024, step=8, value=400, label="Width")
+            h = gr.Slider(64, 1024, step=8, value=400, label="Height")
+        with gr.Row():
+            se = gr.Number(label="Seed", value=0)
+            st = gr.Slider(1, 100, value=10, label="Steps")
+            cf = gr.Slider(1, 20, step=0.1, value=3, label="CFG Scale")
+        with gr.Row():
+            samp = gr.Dropdown(["uni_pc", "euler", "dpmpp_2m", "ddim", "lms"], value="uni_pc", label="Sampler")
+            sched = gr.Dropdown(["simple", "normal", "karras", "exponential"], value="normal", label="Scheduler")
+        with gr.Row():
+            fr = gr.Slider(1, 60, value=2, label="Frames")
+            fps = gr.Slider(1, 60, value=10, label="FPS")
+            fmt = gr.Radio(["mp4", "webm"], value="webm", label="Output Format")
+            q6 = gr.Checkbox(label="Use Q6 UNet model", value=False)
+        gen_btn = gr.Button("Generate")
+        gen_log = gr.Textbox(lines=10, interactive=False, label="Log")
+        gen_out = gr.Video(label="Output Video/Image")
+
+        gen_btn.click(
+            fn=generate_video,
+            inputs=[pos, neg, w, h, se, st, cf, samp, sched, fr, fps, fmt, q6],
+            outputs=[gen_log, gen_out]
+        )
+
+if __name__ == "__main__":
+    app.launch()