app.py

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

# 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

# -------- Helper function to download a file using requests --------
def download_file(url, dest_path):
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
    if os.path.exists(dest_path):
        return f"File already exists: {dest_path}"
    with requests.get(url, stream=True) as r:
        r.raise_for_status()
        with open(dest_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
    return f"Downloaded {url} to {dest_path}"

# -------------------------
# 1. Environment Setup (without aria2c)
# -------------------------
def environment_setup(use_q6: bool):
    global useQ6
    useQ6 = use_q6
    output = []

    # Install Python packages
    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",
        "pip install gradio==5.38.0 imageio numpy Pillow requests -q"
    ]
    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)

    # Clone ComfyUI if missing
    if not os.path.isdir("./ComfyUI"):
        output.append("Cloning ComfyUI repo...")
        proc = subprocess.run("git clone https://github.com/Isi-dev/ComfyUI ./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("./ComfyUI/custom_nodes/ComfyUI_GGUF"):
        output.append("Cloning ComfyUI_GGUF repo...")
        proc = subprocess.run("cd ./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 for custom nodes
        proc = subprocess.run("pip install -r ./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 = "./ComfyUI/models/unet"
    text_enc_dir = "./ComfyUI/models/text_encoders"
    vae_dir = "./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 UNet model using requests fallback
    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"
    unet_path = os.path.join(model_unet_dir, model_name)
    output.append(download_file(model_url, unet_path))

    # 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"
    te_path = os.path.join(text_enc_dir, "umt5_xxl_fp8_e4m3fn_scaled.safetensors")
    vae_path = os.path.join(vae_dir, "wan_2.1_vae.safetensors")
    output.append(download_file(te_url, te_path))
    output.append(download_file(vae_url, vae_path))

    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

    sys.path.insert(0, './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="./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="./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="./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="./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):
    # Return path for Gradio video component
    return video_path

# -------------------------
# 4. Example Gradio interface setup (simplified)
# -------------------------
def dummy_inference(prompt):
    # Placeholder for inference logic
    return f"Prompt received: {prompt}"

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# ComfyUI Integration Demo")

        use_q6_checkbox = gr.Checkbox(label="Use Q6 Model", value=False)
        setup_button = gr.Button("Setup Environment & Download Models")
        setup_output = gr.Textbox(label="Setup Log", lines=15)

        prompt_input = gr.Textbox(label="Prompt")
        run_button = gr.Button("Run Inference")
        result_output = gr.Textbox(label="Output")

        setup_button.click(fn=environment_setup, inputs=[use_q6_checkbox], outputs=[setup_output])
        run_button.click(fn=dummy_inference, inputs=[prompt_input], outputs=[result_output])

    demo.launch()

if __name__ == "__main__":
    main()