Spaces:
Running
Running
| 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) |