Spaces:
Running
Running
| from pydantic import BaseModel, Field | |
| from typing import Optional, Dict, Any | |
| import json, uuid, time, os | |
| import requests | |
| import websocket | |
| from urllib.parse import urlencode | |
| import gradio as gr | |
| COMFY_HOST = os.getenv("COMFY_HOST", "134.199.132.159") | |
| with open("workflow.json", "r", encoding="utf-8") as f: | |
| WORKFLOW_TEMPLATE: Dict[str, Any] = json.load(f) | |
| class T2VReq(BaseModel): | |
| token: str = Field(...) | |
| text: str = Field(...) | |
| negative: Optional[str] = None | |
| seed: Optional[int] = None | |
| steps: Optional[int] = 20 | |
| cfg: Optional[float] = 5.0 | |
| width: Optional[int] = 1280 | |
| height: Optional[int] = 704 | |
| length: Optional[int] = 121 | |
| fps: Optional[int] = 24 | |
| filename_prefix: Optional[str] = "video/ComfyUI" | |
| def _inject_params(prompt: Dict[str, Any], r: T2VReq) -> Dict[str, Any]: | |
| p = json.loads(json.dumps(prompt)) | |
| p["6"]["inputs"]["text"] = r.text | |
| if r.seed is None: | |
| r.seed = int.from_bytes(os.urandom(8), "big") & ((1 << 63) - 1) | |
| p["3"]["inputs"]["seed"] = r.seed | |
| if r.steps is not None: p["3"]["inputs"]["steps"] = r.steps | |
| if r.cfg is not None: p["3"]["inputs"]["cfg"] = r.cfg | |
| if r.width is not None: p["55"]["inputs"]["width"] = r.width | |
| if r.height is not None: p["55"]["inputs"]["height"] = r.height | |
| if r.length is not None: p["55"]["inputs"]["length"] = r.length | |
| if r.fps is not None: p["57"]["inputs"]["fps"] = r.fps | |
| if r.filename_prefix: | |
| p["58"]["inputs"]["filename_prefix"] = r.filename_prefix | |
| return p | |
| def _open_ws(client_id: str, token: str): | |
| ws = websocket.WebSocket() | |
| ws.connect(f"ws://{COMFY_HOST}/ws?clientId={client_id}&token={token}", timeout=1800) | |
| return ws | |
| def _queue_prompt(prompt: Dict[str, Any], client_id: str, token: str) -> str: | |
| payload = {"prompt": prompt, "client_id": client_id} | |
| resp = requests.post(f"http://{COMFY_HOST}/prompt?token={token}", json=payload, timeout=1800) | |
| if resp.status_code != 200: | |
| raise RuntimeError(f"ComfyUI /prompt err: {resp.text}") | |
| data = resp.json() | |
| if "prompt_id" not in data: | |
| raise RuntimeError(f"/prompt no prompt_id: {data}") | |
| return data["prompt_id"] | |
| def _get_history(prompt_id: str, token: str) -> Dict[str, Any]: | |
| r = requests.get(f"http://{COMFY_HOST}/history/{prompt_id}?token={token}", timeout=1800) | |
| r.raise_for_status() | |
| hist = r.json() | |
| return hist.get(prompt_id, {}) | |
| def _extract_video_from_history(history: Dict[str, Any]) -> Dict[str, str]: | |
| outputs = history.get("outputs", {}) | |
| for _, node_out in outputs.items(): | |
| if "images" in node_out: | |
| for it in node_out["images"]: | |
| if all(k in it for k in ("filename", "subfolder", "type")): | |
| fn = it["filename"] | |
| if fn.lower().endswith((".mp4", ".webm", ".gif", ".mov", ".mkv")): | |
| return {"filename": it["filename"], "subfolder": it["subfolder"], "type": it["type"]} | |
| for key in ("videos", "files"): | |
| if key in node_out and node_out[key]: | |
| it = node_out[key][0] | |
| if all(k in it for k in ("filename", "subfolder", "type")): | |
| return {"filename": it["filename"], "subfolder": it["subfolder"], "type": it["type"]} | |
| raise RuntimeError("No video file found in history outputs") | |
| sample_prompts = [ | |
| "A golden retriever running across a beach at sunset, cinematic, 24fps", | |
| "A cyberpunk city street at night with neon lights, light rain, slow pan", | |
| "Aerial drone shot over snowy mountains, 5 seconds, dramatic lighting", | |
| "Cartoon-style cat riding a skateboard in a park, vibrant colors" | |
| ] | |
| with gr.Blocks( | |
| title="T2V UI", | |
| theme=gr.themes.Soft(primary_hue="blue", secondary_hue="blue", neutral_hue="slate"), | |
| ) as demo: | |
| st_token = gr.State() | |
| gr.Markdown("# Wan2.2 T2V running on AMD MI300x") | |
| gr.Markdown("### Prompt") | |
| text = gr.Textbox(label="Prompt", placeholder="Describe the video you want", lines=3) | |
| gr.Examples(examples=sample_prompts, inputs=text) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| with gr.Row(): | |
| width = gr.Number(label="Width", value=1280, precision=0) | |
| height = gr.Number(label="Height", value=704, precision=0) | |
| with gr.Row(): | |
| length = gr.Number(label="Frames", value=121, precision=0) | |
| fps = gr.Number(label="FPS", value=24, precision=0) | |
| with gr.Row(): | |
| steps = gr.Number(label="Steps", value=20, precision=0) | |
| cfg = gr.Number(label="CFG", value=5.0) | |
| seed = gr.Number(label="Seed (optional)", value=None) | |
| filename_prefix = gr.Textbox(label="Filename prefix", value="video/ComfyUI") | |
| run_btn = gr.Button("Generate", variant="primary") | |
| prog_bar = gr.Slider(label="Progress", minimum=0, maximum=100, value=0, step=1, interactive=False) | |
| out_video = gr.Video(label="Result") | |
| def _init_token(): | |
| return str(uuid.uuid4()) | |
| demo.load(_init_token, outputs=st_token) | |
| def generate_fn(text, width, height, length, fps, steps, cfg, seed, filename_prefix, token): | |
| req = T2VReq( | |
| token=token, | |
| text=text, | |
| seed=int(seed) if seed is not None else None, | |
| steps=int(steps) if steps is not None else None, | |
| cfg=float(cfg) if cfg is not None else None, | |
| width=int(width) if width is not None else None, | |
| height=int(height) if height is not None else None, | |
| length=int(length) if length is not None else None, | |
| fps=int(fps) if fps is not None else None, | |
| filename_prefix=filename_prefix if filename_prefix else None, | |
| ) | |
| prompt = _inject_params(WORKFLOW_TEMPLATE, req) | |
| client_id = str(uuid.uuid4()) | |
| ws = _open_ws(client_id, req.token) | |
| prompt_id = _queue_prompt(prompt, client_id, req.token) | |
| total_nodes = max(1, len(prompt)) | |
| seen = set() | |
| p = 0 | |
| last_emit = -1 | |
| start = time.time() | |
| ws.settimeout(60) | |
| while True: | |
| out = ws.recv() | |
| if isinstance(out, (bytes, bytearray)): | |
| if p < 95 and time.time() - start > 2: | |
| p = min(95, p + 1) | |
| if p != last_emit: | |
| last_emit = p | |
| yield p, None | |
| continue | |
| msg = json.loads(out) | |
| if msg.get("type") == "executing": | |
| data = msg.get("data", {}) | |
| if data.get("prompt_id") != prompt_id: | |
| continue | |
| node = data.get("node") | |
| if node is None: | |
| break | |
| if node not in seen: | |
| seen.add(node) | |
| p = min(99, int(len(seen) / total_nodes * 100)) | |
| if p != last_emit: | |
| last_emit = p | |
| yield p, None | |
| ws.close() | |
| hist = _get_history(prompt_id, req.token) | |
| info = _extract_video_from_history(hist) | |
| q = urlencode(info) | |
| video_url = f"http://{COMFY_HOST}/view?{q}&token={req.token}" | |
| yield 100, video_url | |
| run_btn.click( | |
| generate_fn, | |
| inputs=[text, width, height, length, fps, steps, cfg, seed, filename_prefix, st_token], | |
| outputs=[prog_bar, out_video] | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(server_name="0.0.0.0", server_port=9001) |