Update app.py

app.py

@@ -1,28 +1,16 @@
 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 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
+from base64 import b64encode
+import requests
+
+# Globals for model loaders and flags
 unet_loader = None
 clip_loader = None
 clip_encode_positive = None
@@ -31,21 +19,114 @@ 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
+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("/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 for custom nodes
+        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 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
+    global vae_loader, empty_latent_video, ksampler, vae_decode, save_webp, save_webm
+
+    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()
@@ -55,172 +136,87 @@ def imports_initialization():
     empty_latent_video = EmptyHunyuanLatentVideo()
     ksampler = KSampler()
     vae_decode = VAEDecode()
+    save_webp = SaveAnimatedWEBP()
+    save_webm = SaveWEBM()
 
-    return "✅ Imports done and models initialized."
+    return "Imports done and models initialized."
 
-
-# Clean GPU memory
+# -------------------------
+# 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()):
-        try:
-            if torch.is_tensor(obj) or (hasattr(obj, "data") and torch.is_tensor(obj.data)):
-                del obj
-        except:
-            pass
+        if torch.is_tensor(obj) or (hasattr(obj, "data") and torch.is_tensor(obj.data)):
+            del obj
     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]
-        )
+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):
+    # 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__":
-    app.launch()
+    main()