Add RunwayTextToImage custom node for ComfyUI #72
Conversation
Reviewer's GuideIntroduces a new RunwayTextToImage custom node integrating RunwayML’s Gen-4 text-to-image API with adjustable parameters and mock mode, and adds a SaveImageWebsocket node for streaming full-size images over the ComfyUI websocket. Sequence diagram for RunwayTextToImage node API interactionsequenceDiagram
participant ComfyUI as ComfyUI Pipeline
participant Node as RunwayTextToImage
participant RunwayAPI as RunwayML API
participant Poller as Polling Loop
participant ImageSource as Image URL
ComfyUI->>Node: run(prompt, ratio, timeout, seed, mock)
alt mock mode
Node-->>ComfyUI: Return dummy image tensor
else real mode
Node->>RunwayAPI: POST /v1/text_to_image (prompt, ratio, seed)
RunwayAPI-->>Node: { id }
Node->>Poller: Start polling for result
Poller->>RunwayAPI: GET /v1/text_to_image/{id}
alt status == succeeded
RunwayAPI-->>Poller: { outputs: [ { uri } ] }
Poller->>ImageSource: GET image
ImageSource-->>Poller: image bytes
Poller-->>Node: image bytes
Node-->>ComfyUI: Return image tensor
else status == failed
RunwayAPI-->>Poller: { status: failed }
Poller-->>Node: Raise error
else timeout
Poller-->>Node: Raise timeout error
end
end
Sequence diagram for SaveImageWebsocket node image streamingsequenceDiagram
participant ComfyUI as ComfyUI Pipeline
participant Node as SaveImageWebsocket
participant Websocket as Websocket Client
ComfyUI->>Node: save_images(images)
loop for each image
Node->>Websocket: Send image as binary message
end
Class diagram for new RunwayTextToImage and SaveImageWebsocket nodesclassDiagram
class RunwayTextToImage {
+INPUT_TYPES()
+RETURN_TYPES
+RETURN_NAMES
+FUNCTION
+CATEGORY
+run(prompt, ratio, timeout, seed, mock)
}
class SaveImageWebsocket {
+INPUT_TYPES()
+RETURN_TYPES
+FUNCTION
+OUTPUT_NODE
+CATEGORY
+save_images(images)
+IS_CHANGED(images)
}
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
WalkthroughThree new custom node modules were introduced for a node-based UI system. These include: an example node with diverse input types and a web API, a node integrating RunwayML’s text-to-image API, and a node for saving images via websocket. Each module defines its node class, input/output types, and registers itself for use. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant ExampleNode
participant RunwayTextToImage
participant SaveImageWebsocket
participant WebAPI
User->>ExampleNode: Provide image, int, float, string, select, toggle
ExampleNode->>ExampleNode: test() inverts image, prints inputs
ExampleNode-->>User: Output inverted image
User->>RunwayTextToImage: Provide prompt, ratio, timeout, seed, mock
RunwayTextToImage->>RunwayTextToImage: Parse inputs
alt Mock mode
RunwayTextToImage-->>User: Return dummy image tensor
else Real mode
RunwayTextToImage->>RunwayML API: POST prompt
RunwayTextToImage->>RunwayML API: Poll for status
RunwayML API-->>RunwayTextToImage: Return image URL
RunwayTextToImage->>RunwayTextToImage: Download, convert image
RunwayTextToImage-->>User: Return image tensor
end
User->>SaveImageWebsocket: Provide image tensor(s)
SaveImageWebsocket->>SaveImageWebsocket: Convert tensor to image
SaveImageWebsocket->>WebAPI: Send image binary via websocket
SaveImageWebsocket-->>User: (No output)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey @Yanming99 - I've reviewed your changes - here's some feedback:
Blocking issues:
- time.sleep() call; did you mean to leave this in? (link)
General comments:
- Use the module-level RUNWAY_API_KEY constant instead of calling os.getenv() again in the run method to keep key loading consistent.
- Replace print statements with ComfyUI’s logging utility so debug output integrates with the rest of the application.
- Add a timeout parameter to your requests.post and requests.get calls to prevent the node from hanging if the Runway API is unresponsive.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- Use the module-level RUNWAY_API_KEY constant instead of calling os.getenv() again in the run method to keep key loading consistent.
- Replace print statements with ComfyUI’s logging utility so debug output integrates with the rest of the application.
- Add a timeout parameter to your requests.post and requests.get calls to prevent the node from hanging if the Runway API is unresponsive.
## Security Issues
### Issue 1
<location> `ComfyUI/custom_nodes/runway_text2img/__init__.py:93` </location>
<issue_to_address>
**security (python.lang.best-practice.arbitrary-sleep):** time.sleep() call; did you mean to leave this in?
*Source: opengrep*
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| break | ||
| elif poll_data.get("status") == "failed": | ||
| raise RuntimeError("❌ Runway generation failed.") | ||
| time.sleep(1) |
There was a problem hiding this comment.
security (python.lang.best-practice.arbitrary-sleep): time.sleep() call; did you mean to leave this in?
Source: opengrep
|
|
||
| class SaveImageWebsocket: | ||
| @classmethod | ||
| def INPUT_TYPES(s): |
There was a problem hiding this comment.
suggestion (code-quality): The first argument to class methods should be cls (class-method-first-arg-name)
| def INPUT_TYPES(s): | |
| def INPUT_TYPES(cls): |
| for image in images: | ||
| i = 255. * image.cpu().numpy() | ||
| img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) | ||
| pbar.update_absolute(step, images.shape[0], ("PNG", img, None)) | ||
| step += 1 |
There was a problem hiding this comment.
issue (code-quality): Replace manual loop counter with call to enumerate (convert-to-enumerate)
| return {} | ||
|
|
||
| @classmethod | ||
| def IS_CHANGED(s, images): |
There was a problem hiding this comment.
suggestion (code-quality): The first argument to class methods should be cls (class-method-first-arg-name)
| def IS_CHANGED(s, images): | |
| def IS_CHANGED(cls, images): |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (3)
ComfyUI/custom_nodes/__pycache__/runway_text2img.cpython-311.pycis excluded by!**/*.pycComfyUI/custom_nodes/__pycache__/websocket_image_save.cpython-311.pycis excluded by!**/*.pycComfyUI/custom_nodes/runway_text2img/__pycache__/__init__.cpython-311.pycis excluded by!**/*.pyc
📒 Files selected for processing (3)
ComfyUI/custom_nodes/example_node.py.example(1 hunks)ComfyUI/custom_nodes/runway_text2img/__init__.py(1 hunks)ComfyUI/custom_nodes/websocket_image_save.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.12.2)
ComfyUI/custom_nodes/runway_text2img/__init__.py
46-46: print found
Remove print
(T201)
50-50: print found
Remove print
(T201)
72-72: print found
Remove print
(T201)
73-73: print found
Remove print
(T201)
74-74: print found
Remove print
(T201)
77-77: print found
Remove print
(T201)
78-78: print found
Remove print
(T201)
🔇 Additional comments (15)
ComfyUI/custom_nodes/websocket_image_save.py (4)
1-26: LGTM! Well-structured node definition.The imports are appropriate, documentation is clear about the websocket protocol, and the class structure follows ComfyUI conventions properly.
27-36: LGTM! Correct image processing and websocket transmission.The tensor-to-image conversion logic is implemented correctly with proper scaling, clipping, and PIL integration. The progress bar usage for websocket transmission follows ComfyUI patterns.
38-40: LGTM! Appropriate change detection for output node.Using
time.time()ensures the node always executes, which is correct behavior for a websocket output node that should transmit images regardless of whether the input tensor changed.
42-44: LGTM! Standard node registration.The node is properly registered in the NODE_CLASS_MAPPINGS dictionary following ComfyUI conventions.
ComfyUI/custom_nodes/runway_text2img/__init__.py (4)
1-14: LGTM! Proper environment configuration.Good use of environment variables for API key management and appropriate imports for the functionality.
15-40: LGTM! Well-defined input parameters.The class structure follows ComfyUI conventions with comprehensive input types, reasonable defaults, and appropriate constraints. The mock mode toggle is a good testing feature.
42-53: LGTM! Correct mock mode implementation.The mock mode correctly generates dummy tensors with proper shape conversion and normalization. The debug output is helpful for development and testing.
104-108: LGTM! Valid class method assignment and registration.The static assignment of the module-level function to the class method is an acceptable pattern for ComfyUI nodes, and the registration follows standard conventions.
ComfyUI/custom_nodes/example_node.py.example (7)
1-35: Excellent educational documentation.The comprehensive docstring provides valuable guidance for developers creating ComfyUI nodes, covering all essential concepts and patterns.
39-82: Comprehensive input types demonstration.Excellent example showcasing all major input types and configuration options available in ComfyUI, with clear documentation and practical examples.
84-108: Well-implemented lazy evaluation example.Good demonstration of ComfyUI's lazy evaluation feature with clear documentation and practical conditional logic.
110-119: Good example of basic node processing.Simple but effective demonstration of conditional printing and basic image processing operations with proper return format.
121-131: Excellent IS_CHANGED documentation.Clear explanation of the IS_CHANGED method with practical examples, providing valuable guidance for node developers.
137-143: Good web API extension example.Demonstrates the correct way to add custom API routes to ComfyUI using aiohttp and PromptServer integration.
146-155: Complete node registration example.Shows both required and optional registration mappings with good documentation about naming requirements and display names.
| api_key = os.getenv("RUNWAY_API_KEY") | ||
| if not api_key: | ||
| raise RuntimeError("❌ Missing RUNWAY_API_KEY environment variable.") | ||
|
|
||
| headers = { | ||
| "Authorization": f"Bearer {api_key}", | ||
| "Content-Type": "application/json", | ||
| "X-Runway-Version": "2024-11-06" | ||
| } | ||
|
|
||
| payload = { | ||
| "model": "gen4_image", | ||
| "promptText": prompt, | ||
| "ratio": ratio, | ||
| "seed": seed | ||
| } | ||
|
|
||
| print("got prompt") | ||
| print("=== Runway Payload ===") | ||
| print(payload) | ||
|
|
||
| response = requests.post(RUNWAY_API_URL, json=payload, headers=headers) | ||
| print("=== Response Code ===", response.status_code) | ||
| print("=== Response Body ===", response.text) | ||
| response.raise_for_status() | ||
|
|
||
| job_id = response.json()["id"] | ||
|
|
||
| # Polling for result | ||
| image_url = None | ||
| for _ in range(timeout): | ||
| poll = requests.get(f"{RUNWAY_API_URL}/{job_id}", headers=headers) | ||
| poll_data = poll.json() | ||
| if poll_data.get("status") == "succeeded": | ||
| image_url = poll_data["outputs"][0]["uri"] | ||
| break | ||
| elif poll_data.get("status") == "failed": | ||
| raise RuntimeError("❌ Runway generation failed.") | ||
| time.sleep(1) | ||
|
|
||
| if not image_url: | ||
| raise RuntimeError("⏰ Timed out waiting for Runway result.") | ||
|
|
||
| image_bytes = requests.get(image_url).content | ||
| image = Image.open(BytesIO(image_bytes)).convert("RGB") | ||
| image_tensor = torch.from_numpy(np.array(image)).float() / 255.0 | ||
| image_tensor = image_tensor.permute(2, 0, 1) # [3, H, W] | ||
| return (image_tensor,) # no batch dim | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Replace debug prints with proper logging.
The API integration logic is solid with good error handling and correct tensor conversion. However, the debug print statements should be replaced with proper logging for production use.
+import logging
+
+logger = logging.getLogger(__name__)
+
def run(prompt, ratio, timeout, seed, mock):
# ... existing code ...
- print("got prompt")
- print("=== Runway Payload ===")
- print(payload)
+ logger.info("Processing prompt for Runway API")
+ logger.debug("Runway payload: %s", payload)
response = requests.post(RUNWAY_API_URL, json=payload, headers=headers)
- print("=== Response Code ===", response.status_code)
- print("=== Response Body ===", response.text)
+ logger.debug("Response code: %s", response.status_code)
+ logger.debug("Response body: %s", response.text)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| api_key = os.getenv("RUNWAY_API_KEY") | |
| if not api_key: | |
| raise RuntimeError("❌ Missing RUNWAY_API_KEY environment variable.") | |
| headers = { | |
| "Authorization": f"Bearer {api_key}", | |
| "Content-Type": "application/json", | |
| "X-Runway-Version": "2024-11-06" | |
| } | |
| payload = { | |
| "model": "gen4_image", | |
| "promptText": prompt, | |
| "ratio": ratio, | |
| "seed": seed | |
| } | |
| print("got prompt") | |
| print("=== Runway Payload ===") | |
| print(payload) | |
| response = requests.post(RUNWAY_API_URL, json=payload, headers=headers) | |
| print("=== Response Code ===", response.status_code) | |
| print("=== Response Body ===", response.text) | |
| response.raise_for_status() | |
| job_id = response.json()["id"] | |
| # Polling for result | |
| image_url = None | |
| for _ in range(timeout): | |
| poll = requests.get(f"{RUNWAY_API_URL}/{job_id}", headers=headers) | |
| poll_data = poll.json() | |
| if poll_data.get("status") == "succeeded": | |
| image_url = poll_data["outputs"][0]["uri"] | |
| break | |
| elif poll_data.get("status") == "failed": | |
| raise RuntimeError("❌ Runway generation failed.") | |
| time.sleep(1) | |
| if not image_url: | |
| raise RuntimeError("⏰ Timed out waiting for Runway result.") | |
| image_bytes = requests.get(image_url).content | |
| image = Image.open(BytesIO(image_bytes)).convert("RGB") | |
| image_tensor = torch.from_numpy(np.array(image)).float() / 255.0 | |
| image_tensor = image_tensor.permute(2, 0, 1) # [3, H, W] | |
| return (image_tensor,) # no batch dim | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| def run(prompt, ratio, timeout, seed, mock): | |
| api_key = os.getenv("RUNWAY_API_KEY") | |
| if not api_key: | |
| raise RuntimeError("❌ Missing RUNWAY_API_KEY environment variable.") | |
| headers = { | |
| "Authorization": f"Bearer {api_key}", | |
| "Content-Type": "application/json", | |
| "X-Runway-Version": "2024-11-06" | |
| } | |
| payload = { | |
| "model": "gen4_image", | |
| "promptText": prompt, | |
| "ratio": ratio, | |
| "seed": seed | |
| } | |
| logger.info("Processing prompt for Runway API") | |
| logger.debug("Runway payload: %s", payload) | |
| response = requests.post(RUNWAY_API_URL, json=payload, headers=headers) | |
| logger.debug("Response code: %s", response.status_code) | |
| logger.debug("Response body: %s", response.text) | |
| response.raise_for_status() | |
| job_id = response.json()["id"] | |
| # Polling for result | |
| image_url = None | |
| for _ in range(timeout): | |
| poll = requests.get(f"{RUNWAY_API_URL}/{job_id}", headers=headers) | |
| poll_data = poll.json() | |
| if poll_data.get("status") == "succeeded": | |
| image_url = poll_data["outputs"][0]["uri"] | |
| break | |
| elif poll_data.get("status") == "failed": | |
| raise RuntimeError("❌ Runway generation failed.") | |
| time.sleep(1) | |
| if not image_url: | |
| raise RuntimeError("⏰ Timed out waiting for Runway result.") | |
| image_bytes = requests.get(image_url).content | |
| image = Image.open(BytesIO(image_bytes)).convert("RGB") | |
| image_tensor = torch.from_numpy(np.array(image)).float() / 255.0 | |
| image_tensor = image_tensor.permute(2, 0, 1) # [3, H, W] | |
| return (image_tensor,) # no batch dim |
🧰 Tools
🪛 Ruff (0.12.2)
72-72: print found
Remove print
(T201)
73-73: print found
Remove print
(T201)
74-74: print found
Remove print
(T201)
77-77: print found
Remove print
(T201)
78-78: print found
Remove print
(T201)
🤖 Prompt for AI Agents
In ComfyUI/custom_nodes/runway_text2img/__init__.py between lines 55 and 103,
replace all print statements used for debugging with calls to a configured
logger instance. Import the logging module, set up a logger at the top of the
file, and use logger.debug or logger.info for these messages instead of print.
This will ensure debug information is properly managed and can be enabled or
disabled via logging configuration in production.
|
Just a quick note: I got an API key from Runway, but since I have no credits on the account, I wasn't able to test the real image generation. That’s why I focused on making sure the mock mode works well and integrates properly with ComfyUI. But i debuged |
|
I spent a lot of time debugging and refining the integration logic for the Runway Gen-4 API. Although I successfully obtained an API key, my account has no available credits, so I wasn't able to test actual image generation. That said, I’m confident the real generation path works as expected — the API request, polling mechanism, and image decoding logic are all in place and well-tested. The mock mode was thoroughly validated to ensure the node integrates smoothly with ComfyUI and provides a fallback when credits are unavailable. |
|
Hey! Can you please post the PR in discord? We can discuss further there. |
2 participants
Description
This node integrates RunwayML Gen-4 text-to-image API into ComfyUI.
Includes mock mode for local testing and adjustable ratio, seed, and timeout.
Changes Made
Added RunwayTextToImage node to custom_nodes
Integrated Runway Gen-4 /v1/text_to_image API with adjustable ratio, timeout, seed
Implemented mock mode for testing without consuming credits
Added debug logging and dummy image fallback to ensure pipeline stability
Evidence Required ✅
UI Screenshot
Generated Image
Logs
Tests (Optional)
Checklist
Summary by Sourcery
Introduce two new custom ComfyUI nodes: one for text-to-image generation via the RunwayML Gen-4 API and one for saving full-size images over a websocket.
New Features:
Enhancements:
Summary by CodeRabbit