app.py

import gradio as gr
import subprocess
import os
import sys
import torch
import gc
import numpy as np
from PIL import Image
import imageio
from base64 import b64encode

# Globals for model loaders and flags
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
save_webp = None
save_webm = None
useQ6 = False

# -------------------------
# 1. Environment Setup
# -------------------------
def environment_setup(use_q6: bool):
    global useQ6
    useQ6 = use_q6
    setup_cmds = [
        "pip install torch==2.6.0 torchvision==0.21.0 -q",
        "pip install torchsde einops diffusers accelerate xformers==0.0.29.post2 -q",
        "pip install av -q",
        "apt -y install aria2 ffmpeg -qq"
    ]
    output = []
    for cmd in setup_cmds:
        output.append(f"Running: {cmd}")
        proc = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        output.append(proc.stdout)
        output.append(proc.stderr)

    if not os.path.isdir("/content/ComfyUI"):
        output.append("Cloning ComfyUI repo...")
        proc = subprocess.run("git clone https://github.com/Isi-dev/ComfyUI /content/ComfyUI", shell=True, capture_output=True, text=True)
        output.append(proc.stdout + proc.stderr)
    else:
        output.append("ComfyUI repo already exists")

    # Clone custom nodes repo
    if not os.path.isdir("/content/ComfyUI/custom_nodes/ComfyUI_GGUF"):
        output.append("Cloning ComfyUI_GGUF repo...")
        proc = subprocess.run("cd /content/ComfyUI/custom_nodes && git clone https://github.com/Isi-dev/ComfyUI_GGUF.git", shell=True, capture_output=True, text=True)
        output.append(proc.stdout + proc.stderr)
        # Install requirements
        proc = subprocess.run("pip install -r /content/ComfyUI/custom_nodes/ComfyUI_GGUF/requirements.txt", shell=True, capture_output=True, text=True)
        output.append(proc.stdout + proc.stderr)
    else:
        output.append("ComfyUI_GGUF repo already exists")

    # Ensure model directories exist
    model_unet_dir = "/content/ComfyUI/models/unet"
    text_enc_dir = "/content/ComfyUI/models/text_encoders"
    vae_dir = "/content/ComfyUI/models/vae"
    os.makedirs(model_unet_dir, exist_ok=True)
    os.makedirs(text_enc_dir, exist_ok=True)
    os.makedirs(vae_dir, exist_ok=True)

    # Download models based on useQ6
    if useQ6:
        model_url = "https://huggingface.co/city96/Wan2.1-T2V-14B-gguf/resolve/main/wan2.1-t2v-14b-Q6_K.gguf"
        model_name = "wan2.1-t2v-14b-Q6_K.gguf"
    else:
        model_url = "https://huggingface.co/city96/Wan2.1-T2V-14B-gguf/resolve/main/wan2.1-t2v-14b-Q5_0.gguf"
        model_name = "wan2.1-t2v-14b-Q5_0.gguf"
    aria2_cmd = f"aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {model_url} -d {model_unet_dir} -o {model_name}"
    output.append(f"Downloading UNet model: {model_name}")
    proc = subprocess.run(aria2_cmd, shell=True, capture_output=True, text=True)
    output.append(proc.stdout + proc.stderr)

    # Download text encoder and VAE
    te_url = "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors"
    vae_url = "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/resolve/main/split_files/vae/wan_2.1_vae.safetensors"
    aria2_te_cmd = f"aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {te_url} -d {text_enc_dir} -o umt5_xxl_fp8_e4m3fn_scaled.safetensors"
    aria2_vae_cmd = f"aria2c --console-log-level=error -c -x 16 -s 16 -k 1M {vae_url} -d {vae_dir} -o wan_2.1_vae.safetensors"
    output.append("Downloading text encoder...")
    proc = subprocess.run(aria2_te_cmd, shell=True, capture_output=True, text=True)
    output.append(proc.stdout + proc.stderr)
    output.append("Downloading VAE...")
    proc = subprocess.run(aria2_vae_cmd, shell=True, capture_output=True, text=True)
    output.append(proc.stdout + proc.stderr)
    return "\n".join(output)

# -------------------------
# 2. Imports & Initialization
# -------------------------
def imports_initialization():
    global unet_loader, clip_loader, clip_encode_positive, clip_encode_negative
    global vae_loader, empty_latent_video, ksampler, vae_decode, save_webp, save_webm
    import sys
    sys.path.insert(0, '/content/ComfyUI')
    from comfy import model_management
    from nodes import (
        CheckpointLoaderSimple,
        CLIPLoader,
        CLIPTextEncode,
        VAEDecode,
        VAELoader,
        KSampler,
        UNETLoader
    )
    from custom_nodes.ComfyUI_GGUF.nodes import UnetLoaderGGUF
    from comfy_extras.nodes_model_advanced import ModelSamplingSD3
    from comfy_extras.nodes_hunyuan import EmptyHunyuanLatentVideo
    from comfy_extras.nodes_images import SaveAnimatedWEBP
    from comfy_extras.nodes_video import SaveWEBM

    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()
    save_webp = SaveAnimatedWEBP()
    save_webm = SaveWEBM()
    return "Imports done and models initialized."

# -------------------------
# 3. Utility Functions
# -------------------------
def clear_memory():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    for obj in list(globals().values()):
        if torch.is_tensor(obj) or (hasattr(obj, "data") and torch.is_tensor(obj.data)):
            del obj
    gc.collect()

def save_as_mp4(images, filename_prefix, fps, output_dir="/content/ComfyUI/output"):
    os.makedirs(output_dir, exist_ok=True)
    output_path = f"{output_dir}/{filename_prefix}.mp4"
    frames = [(img.cpu().numpy() * 255).astype(np.uint8) for img in images]
    with imageio.get_writer(output_path, fps=fps) as writer:
        for frame in frames:
            writer.append_data(frame)
    return output_path

def save_as_webp(images, filename_prefix, fps, quality=90, lossless=False, method=4, output_dir="/content/ComfyUI/output"):
    os.makedirs(output_dir, exist_ok=True)
    output_path = f"{output_dir}/{filename_prefix}.webp"
    frames = [(img.cpu().numpy() * 255).astype(np.uint8) for img in images]
    kwargs = {'fps': int(fps), 'quality': int(quality), 'lossless': bool(lossless), 'method': int(method)}
    with imageio.get_writer(output_path, format='WEBP', mode='I', **kwargs) as writer:
        for frame in frames:
            writer.append_data(frame)
    return output_path

def save_as_webm(images, filename_prefix, fps, codec="vp9", quality=32, output_dir="/content/ComfyUI/output"):
    os.makedirs(output_dir, exist_ok=True)
    output_path = f"{output_dir}/{filename_prefix}.webm"
    frames = [(img.cpu().numpy() * 255).astype(np.uint8) for img in images]
    kwargs = {'fps': int(fps), 'quality': int(quality), 'codec': str(codec), 'output_params': ['-crf', str(int(quality))]}
    with imageio.get_writer(output_path, format='FFMPEG', mode='I', **kwargs) as writer:
        for frame in frames:
            writer.append_data(frame)
    return output_path

def save_as_image(image, filename_prefix, output_dir="/content/ComfyUI/output"):
    os.makedirs(output_dir, exist_ok=True)
    output_path = f"{output_dir}/{filename_prefix}.png"
    frame = (image.cpu().numpy() * 255).astype(np.uint8)
    Image.fromarray(frame).save(output_path)
    return output_path

def display_video_gradio(video_path):
    # Returns video path for Gradio video component, no HTML needed.
    return video_path

# -------------------------
# 4. Generation Function
# -------------------------
def generate_video(
    positive_prompt,
    negative_prompt,
    width,
    height,
    seed,
    steps,
    cfg_scale,
    sampler_name,
    scheduler,
    frames,
    fps,
    output_format
):
    global useQ6
    with torch.inference_mode():
        log = []
        log.append("Loading Text_Encoder...")
        clip = clip_loader.load_clip("umt5_xxl_fp8_e4m3fn_scaled.safetensors", "wan", "default")[0]
        positive = clip_encode_positive.encode(clip, positive_prompt)[0]
        negative = clip_encode_negative.encode(clip, negative_prompt)[0]
        del clip
        torch.cuda.empty_cache()
        gc.collect()

        empty_latent = empty_latent_video.generate(width, height, frames, 1)[0]

        log.append("Loading Unet Model...")
        if useQ6:
            model = unet_loader.load_unet("wan2.1-t2v-14b-Q6_K.gguf")[0]
        else:
            model = unet_loader.load_unet("wan2.1-t2v-14b-Q5_0.gguf")[0]

        log.append("Generating video...")
        sampled = ksampler.sample(
            model=model,
            seed=seed,
            steps=steps,
            cfg=cfg_scale,
            sampler_name=sampler_name,
            scheduler=scheduler,
            positive=positive,
            negative=negative,
            latent_image=empty_latent
        )[0]

        del model
        torch.cuda.empty_cache()
        gc.collect()

        log.append("Loading VAE...")
        vae = vae_loader.load_vae("wan_2.1_vae.safetensors")[0]

        output_path = ""
        try:
            log.append("Decoding latents...")
            decoded = vae_decode.decode(vae, sampled)[0]
            del vae
            torch.cuda.empty_cache()
            gc.collect()

            if frames == 1:
                log.append("Single frame - saving as PNG image...")
                output_path = save_as_image(decoded[0], "ComfyUI")
            else:
                if output_format.lower() == "webm":
                    log.append("Saving as WEBM...")
                    output_path = save_as_webm(decoded, "ComfyUI", fps=fps, codec="vp9", quality=10)
                elif output_format.lower() == "mp4":
                    log.append("Saving as MP4...")
                    output_path = save_as_mp4(decoded, "ComfyUI", fps)
                else:
                    log.append(f"Unsupported output format: {output_format}")
                    return "\n".join(log), None

        except Exception as e:
            log.append(f"Error: {str(e)}")
            return "\n".join(log), None
        finally:
            clear_memory()

        return "\n".join(log), output_path

# -------------------------
# 5. Run Example (with default/custom params)
# -------------------------
def generate_video_example():
    params = {
        'positive_prompt': "lion",
        'negative_prompt': "Bright tones, overexposure, static, blurry details, subtitles, artistic style, artwork, painting, still image, dull overall, worst quality, low quality, JPEG compression artifacts, ugly, deformed, extra fingers, poorly drawn hands, poorly drawn face, disfigured, malformed limbs, finger fusion, static frame, messy background, three legs, crowded background, walking backwards",
        'width': 400,
        'height': 400,
        'seed': 82628696717258,
        'steps': 10,
        'cfg_scale': 3,
        'sampler_name': "uni_pc",
        'scheduler': "normal",
        'frames': 2,
        'fps': 10,
        'output_format': "webm",
    }
    return generate_video(**params)

# -----------------------------------
# GRADIO INTERFACE SETUP
# -----------------------------------

with gr.Blocks() as demo:
    gr.Markdown("# ComfyUI Video/Image Generation with Gradio (Colab T4 GPU)")

    with gr.Tab("1. Environment Setup"):
        q6_checkbox = gr.Checkbox(label="Use Q6 model (else Q5)", value=False)
        env_out = gr.Textbox(label="Setup Log", lines=15, interactive=False)
        env_button = gr.Button("Run Environment Setup")
        env_button.click(fn=environment_setup, inputs=q6_checkbox, outputs=env_out)

    with gr.Tab("2. Imports & Initialization"):
        init_button = gr.Button("Initialize Imports and Models")
        init_out = gr.Textbox(label="Initialization Status", interactive=False)
        init_button.click(fn=imports_initialization, inputs=None, outputs=init_out)

    with gr.Tab("4. Generate Video/Image"):
        with gr.Row():
            pos_prompt = gr.Textbox(label="Positive Prompt", value="lion")
            neg_prompt = gr.Textbox(label="Negative Prompt", value="Bright tones, overexposure, static, blurry details, subtitles, artistic style, artwork, painting, still image, dull overall, worst quality, low quality, JPEG compression artifacts, ugly, deformed, extra fingers, poorly drawn hands, poorly drawn face, disfigured, malformed limbs, finger fusion, static frame, messy background, three legs, crowded background, walking backwards")
        with gr.Row():
            width_slider = gr.Slider(64, 1024, value=400, step=8, label="Width")
            height_slider = gr.Slider(64, 1024, value=400, step=8, label="Height")
        with gr.Row():
            seed_num = gr.Number(value=82628696717258, label="Seed")
            steps_slider = gr.Slider(1, 100, value=10, step=1, label="Steps")
            cfg_slider = gr.Slider(1, 20, value=3, step=0.1, label="CFG Scale")
        with gr.Row():
            sampler_dropdown = gr.Dropdown(choices=["uni_pc", "euler", "dpmpp_2m", "ddim", "lms"], value="uni_pc", label="Sampler")
            scheduler_dropdown = gr.Dropdown(choices=["simple", "normal", "karras", "exponential"], value="normal", label="Scheduler")
        with gr.Row():
            frames_slider = gr.Slider(1, 120, value=2, step=1, label="Frames")
            fps_slider = gr.Slider(1, 60, value=10, step=1, label="FPS")
            output_format_radio = gr.Radio(choices=["mp4", "webm"], value="webm", label="Output Format")

        gen_button = gr.Button("Generate")
        gen_log = gr.Textbox(label="Generation Log", lines=15, interactive=False)
        gen_video = gr.Video(label="Output Video/Image")

        gen_button.click(
            fn=generate_video,
            inputs=[
                pos_prompt, neg_prompt, width_slider, height_slider,
                seed_num, steps_slider, cfg_slider, sampler_dropdown,
                scheduler_dropdown, frames_slider, fps_slider, output_format_radio
            ],
            outputs=[gen_log, gen_video]
        )

    with gr.Tab("5. Run Example"):
        example_button = gr.Button("Run Example Generation")
        example_log = gr.Textbox(label="Example Run Log", lines=15, interactive=False)
        example_video = gr.Video(label="Example Output")

        example_button.click(
            fn=generate_video_example,
            inputs=None,
            outputs=[example_log, example_video]
        )

    with gr.Tab("3. Utility Functions"):
        gr.Markdown(
            """
            Utility functions like clearing memory and saving files are used internally in this app.
            """)

# Launch app
if __name__ == "__main__":
    demo.launch(share=True)