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