Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
df083ed8be | ||
|
|
641f656ed5 | ||
|
|
bdf71154fc | ||
|
|
325e2610a4 | ||
|
|
e10b421b23 | ||
|
|
3e98047524 | ||
|
|
bf35dc72c6 | ||
|
|
bb8166cc8f | ||
|
|
9ca5218ce0 | ||
|
|
c498711f10 | ||
|
|
6de50104f1 | ||
|
|
348599a139 | ||
|
|
52004b753d | ||
|
|
519e6025a8 | ||
|
|
2c4277711a | ||
|
|
195842c634 | ||
|
|
851b2ef353 | ||
|
|
0b723ce1df | ||
|
|
29d2c9856a | ||
|
|
9520576e4e | ||
|
|
13d2114df9 | ||
|
|
fc4a3d37cd | ||
|
|
37a444993e | ||
|
|
128e19a9b5 | ||
|
|
d15f6ffc8c | ||
|
|
510279d4d7 | ||
|
|
1697364e13 | ||
|
|
ed2d1388ec | ||
|
|
de5d753cdf | ||
|
|
523fe68f6b | ||
|
|
ad360061d7 | ||
|
|
98bbda1f94 | ||
|
|
c4a96e3484 | ||
|
|
0e655a4303 | ||
|
|
ec09d32cf6 | ||
|
|
9c8e73e578 | ||
|
|
af9936a706 | ||
|
|
9716a63591 | ||
|
|
7775d0b212 | ||
|
|
59c04a6b84 | ||
|
|
6c4c2ebc94 | ||
|
|
89def34b23 | ||
|
|
411336a217 | ||
|
|
8a2793e3fa | ||
|
|
194654ff1a | ||
|
|
fa58cddf73 | ||
|
|
4b2a99e5f5 | ||
|
|
766168ec1b | ||
|
|
87112b8b2a | ||
|
|
73b2bfe5b1 | ||
|
|
b3d8f9d16f | ||
|
|
b63dac6cad | ||
|
|
3ee819d042 | ||
|
|
1ba7122694 | ||
|
|
e00194e0e7 | ||
|
|
c74b572c38 | ||
|
|
5d23e30dbc | ||
|
|
7382295a76 | ||
|
|
c4ef38cff3 | ||
|
|
c1b1e2497c | ||
|
|
7cf1b45545 | ||
|
|
0c7c1d83be | ||
|
|
f53010ceef | ||
|
|
8206c475e5 | ||
|
|
f2e1779450 | ||
|
|
8ffb72bb4d | ||
|
|
6af9e26255 | ||
|
|
e7fde9951a | ||
|
|
def123c242 | ||
|
|
436b773bd8 | ||
|
|
b083702eb9 | ||
|
|
618b9ae4bb | ||
|
|
569f76f484 | ||
|
|
0d6ed9d9c5 | ||
|
|
cc07d20f86 | ||
|
|
36c5d06cf8 | ||
|
|
b8aadd7faa | ||
|
|
1e94a69e5f | ||
|
|
87c97e04a3 | ||
|
|
76da850b3d | ||
|
|
a208cd482b | ||
|
|
9295546c17 | ||
|
|
852840dcae | ||
|
|
06d7b955f5 | ||
|
|
64179a3170 | ||
|
|
1a84f62ec1 | ||
|
|
d7b20fc717 | ||
|
|
95e00c65bf | ||
|
|
4d34e27463 | ||
|
|
5c6f15469d | ||
|
|
9acae29f27 | ||
|
|
83f0059720 | ||
|
|
81010bbb4e | ||
|
|
c4dcfd586d | ||
|
|
e3e95ab197 | ||
|
|
7cc72606a5 | ||
|
|
8be6270aef | ||
|
|
b944db244e | ||
|
|
d99aa60623 | ||
|
|
214ecf35eb | ||
|
|
711aeed35d | ||
|
|
a9ff597d5c | ||
|
|
e1f9b50c3d | ||
|
|
87a0a5d5fe | ||
|
|
45a1627992 | ||
|
|
1c9ab10e16 | ||
|
|
f67a978c5d | ||
|
|
291fa65a7b | ||
|
|
c5e79124cc | ||
|
|
5610f51a71 | ||
|
|
3fd583c483 | ||
|
|
8888c2cb0b | ||
|
|
35178b92c8 | ||
|
|
d35caae023 | ||
|
|
3f553e600a | ||
|
|
cd2052757a | ||
|
|
73cf5fc0dd | ||
|
|
2a37c1b410 | ||
|
|
24d6912396 | ||
|
|
15883ebb66 | ||
|
|
decd7def28 | ||
|
|
491e4a2758 | ||
|
|
a81ae85d1e | ||
|
|
f49f65a196 | ||
|
|
6558997050 | ||
|
|
cc38033f6d | ||
|
|
9806a7f576 | ||
|
|
7ace1668d7 | ||
|
|
b05b89148e | ||
|
|
016a63d600 | ||
|
|
b912148092 | ||
|
|
8b32435303 | ||
|
|
4bc286ef65 | ||
|
|
c7e7f5a6d9 | ||
|
|
e1b70b3db2 | ||
|
|
c025ecade6 | ||
|
|
daad0c78e3 | ||
|
|
1a58f5ca58 | ||
|
|
6fa80e084d | ||
|
|
33d3f4b5b2 | ||
|
|
5c38a81603 | ||
|
|
90ad202495 | ||
|
|
746ac77ab7 | ||
|
|
8446ae8f8f | ||
|
|
80101b3f60 | ||
|
|
52c8a4430c | ||
|
|
af580e1895 | ||
|
|
f74158ecca | ||
|
|
d189d5ba67 | ||
|
|
593b587f36 | ||
|
|
ca459e6f44 | ||
|
|
a128b92e33 | ||
|
|
8696680d87 | ||
|
|
ad26e34c4f | ||
|
|
7d7b304eaf | ||
|
|
49064d08ae | ||
|
|
ae97964de7 | ||
|
|
5539eb2417 | ||
|
|
e241a6817a | ||
|
|
da2e119664 | ||
|
|
0809d28431 | ||
|
|
6c282d3f1b | ||
|
|
deba755ed2 | ||
|
|
e7d8843524 | ||
|
|
6ca9bb15b3 | ||
|
|
2220f76fea | ||
|
|
5755166e41 | ||
|
|
08027b34e7 | ||
|
|
c8838160ec | ||
|
|
b1b7b04ce3 | ||
|
|
2530cfebfc | ||
|
|
8be9e6e379 | ||
|
|
1b2649759e | ||
|
|
29a48ae943 | ||
|
|
2d2e848f65 | ||
|
|
a9235af50e | ||
|
|
981df1c7b3 | ||
|
|
d68bb453bf | ||
|
|
a125a0502f | ||
|
|
b931cb7a62 | ||
|
|
6f86e35cea | ||
|
|
0e6e99b2a0 | ||
|
|
20ad041d5b | ||
|
|
734597eaa9 | ||
|
|
c2fbbaa237 | ||
|
|
5e26c6af1d | ||
|
|
2ed961690a | ||
|
|
db3f6f487f | ||
|
|
9fb2710c06 | ||
|
|
c412807145 | ||
|
|
dafd978f18 | ||
|
|
0eda79eab9 | ||
|
|
3ffaf352b2 | ||
|
|
ca6d9db364 | ||
|
|
f9f7bba167 | ||
|
|
331bdcbc99 | ||
|
|
f730c6cef5 | ||
|
|
d5fba647a8 | ||
|
|
4948d89783 | ||
|
|
6defb15da0 | ||
|
|
b1347275a4 | ||
|
|
8dc1b3fecf | ||
|
|
aac823bfa8 | ||
|
|
28221907a4 | ||
|
|
0f541d334a | ||
|
|
4926cdd018 | ||
|
|
a27038b236 | ||
|
|
c27216b3bf | ||
|
|
17ad02eac4 | ||
|
|
8a7e4af30b | ||
|
|
3ad72f35c6 | ||
|
|
5e68b9145a | ||
|
|
bd179ccf86 | ||
|
|
3b2e21e1d3 | ||
|
|
916c2805cd | ||
|
|
7dbcf3b8f4 | ||
|
|
a2eea1cfa5 | ||
|
|
16ef56460f | ||
|
|
2d99246afe | ||
|
|
31aa0bb62e |
@@ -0,0 +1,48 @@
|
||||
# ComfyUI-DiscordSend Bot Environment Variables
|
||||
# Copy this file to .env and fill in your values
|
||||
# These override settings in config.yaml
|
||||
|
||||
# =============================================================================
|
||||
# REQUIRED
|
||||
# =============================================================================
|
||||
|
||||
# Discord bot token from Discord Developer Portal
|
||||
DISCORDBOT_DISCORD_TOKEN=your_bot_token_here
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL - ComfyUI Connection
|
||||
# =============================================================================
|
||||
|
||||
# ComfyUI server URL (default: http://127.0.0.1:8188)
|
||||
# DISCORDBOT_COMFYUI_URL=http://127.0.0.1:8188
|
||||
|
||||
# ComfyUI WebSocket URL (default: auto-derived from COMFYUI_URL)
|
||||
# DISCORDBOT_COMFYUI_WS_URL=ws://127.0.0.1:8188/ws
|
||||
|
||||
# Request timeout in seconds (default: 30)
|
||||
# DISCORDBOT_COMFYUI_TIMEOUT=30
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL - Database
|
||||
# =============================================================================
|
||||
|
||||
# Database URL (default: SQLite in bot/data/bot.db)
|
||||
# For PostgreSQL: postgresql+asyncpg://user:pass@localhost/dbname
|
||||
# DISCORDBOT_DATABASE_URL=sqlite+aiosqlite:///bot/data/bot.db
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL - Defaults
|
||||
# =============================================================================
|
||||
|
||||
# Maximum pending jobs per user (default: 3)
|
||||
# DISCORDBOT_MAX_QUEUE_PER_USER=3
|
||||
|
||||
# Path to default workflow JSON file
|
||||
# DISCORDBOT_WORKFLOW_PATH=bot/data/default_workflow_api.json
|
||||
|
||||
# =============================================================================
|
||||
# OPTIONAL - Discord Application
|
||||
# =============================================================================
|
||||
|
||||
# Discord application ID (for slash command registration)
|
||||
# DISCORDBOT_APPLICATION_ID=your_app_id_here
|
||||
@@ -0,0 +1,115 @@
|
||||
name: Update Clone Count Badge
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: '0 0 * * *' # Runs daily at midnight
|
||||
workflow_dispatch: # Allows manual trigger
|
||||
|
||||
jobs:
|
||||
update-badge:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write # Needed to push to the branch
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v3
|
||||
|
||||
- name: Fetch Traffic Stats and Save JSON
|
||||
uses: actions/github-script@v6
|
||||
with:
|
||||
github-token: ${{ secrets.TRAFFIC_TOKEN }}
|
||||
script: |
|
||||
try {
|
||||
const { owner, repo } = context.repo;
|
||||
|
||||
// 1. Fetch Clones
|
||||
const clones = await github.rest.repos.getClones({
|
||||
owner,
|
||||
repo,
|
||||
});
|
||||
|
||||
// 2. Fetch Views (Visitors)
|
||||
const views = await github.rest.repos.getViews({
|
||||
owner,
|
||||
repo,
|
||||
});
|
||||
|
||||
// 3. Fetch Releases for Smart Download Count
|
||||
// We need to iterate through all releases to get the total count
|
||||
let smartDownloadCount = 0;
|
||||
let page = 1;
|
||||
let releases = [];
|
||||
|
||||
do {
|
||||
const response = await github.rest.repos.listReleases({
|
||||
owner,
|
||||
repo,
|
||||
per_page: 100,
|
||||
page: page
|
||||
});
|
||||
releases = response.data;
|
||||
|
||||
for (const release of releases) {
|
||||
let maxDownloads = 0;
|
||||
for (const asset of release.assets) {
|
||||
if (asset.download_count > maxDownloads) {
|
||||
maxDownloads = asset.download_count;
|
||||
}
|
||||
}
|
||||
smartDownloadCount += maxDownloads;
|
||||
}
|
||||
page++;
|
||||
} while (releases.length === 100);
|
||||
|
||||
const fs = require('fs');
|
||||
const data = {
|
||||
clones: {
|
||||
count: clones.data.count,
|
||||
uniques: clones.data.uniques,
|
||||
},
|
||||
views: {
|
||||
count: views.data.count,
|
||||
uniques: views.data.uniques,
|
||||
},
|
||||
downloads: {
|
||||
smart_count: smartDownloadCount
|
||||
},
|
||||
timestamp: new Date().toISOString()
|
||||
};
|
||||
|
||||
console.log('Stats fetched:', JSON.stringify(data));
|
||||
|
||||
// Write to a generic stats file
|
||||
fs.writeFileSync('traffic_stats.json', JSON.stringify(data, null, 2));
|
||||
// Maintain the old file for backward compatibility if needed, or just switch everything
|
||||
fs.writeFileSync('git_clones.json', JSON.stringify(data.clones, null, 2));
|
||||
|
||||
} catch (error) {
|
||||
console.error('Error fetching stats:', error);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
- name: Push to Badges Branch
|
||||
run: |
|
||||
git config --global user.name 'github-actions[bot]'
|
||||
git config --global user.email 'github-actions[bot]@users.noreply.github.com'
|
||||
|
||||
# Stash the files
|
||||
mv traffic_stats.json /tmp/traffic_stats.json
|
||||
mv git_clones.json /tmp/git_clones.json
|
||||
|
||||
# Fetch the badges branch or create it
|
||||
git fetch origin badges:badges || true
|
||||
git checkout badges || git checkout --orphan badges
|
||||
|
||||
# Clean the branch to ensure it only has the JSON
|
||||
git rm -rf .
|
||||
|
||||
# Restore the files
|
||||
mv /tmp/traffic_stats.json traffic_stats.json
|
||||
mv /tmp/git_clones.json git_clones.json
|
||||
|
||||
# Commit and push
|
||||
git add traffic_stats.json git_clones.json
|
||||
git diff --quiet && git diff --staged --quiet || (git commit -m "Update traffic statistics" && git push origin badges)
|
||||
@@ -0,0 +1,20 @@
|
||||
name: Publish to Comfy registry
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "pyproject.toml"
|
||||
|
||||
jobs:
|
||||
publish-node:
|
||||
name: Publish Custom Node to registry
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
# Sensitive files - NEVER commit these
|
||||
.env
|
||||
.env.*
|
||||
!.env.example
|
||||
config.yaml
|
||||
config.yml
|
||||
*.pem
|
||||
*.key
|
||||
secrets.*
|
||||
credentials.*
|
||||
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
Errors.md
|
||||
Claude_Last_Convo.md
|
||||
ANALYSIS.md
|
||||
REPO_ANALYSIS.md
|
||||
PROJECT_BREAKDOWN.md
|
||||
REFACTOR_PRD.md
|
||||
PRD.md
|
||||
REPO_BREAKDOWN.md
|
||||
.jules/
|
||||
.Jules/
|
||||
.claude/
|
||||
.venv/
|
||||
venv/
|
||||
|
||||
# Standard Python/Test Artifacts
|
||||
.pytest_cache/
|
||||
htmlcov/
|
||||
.coverage
|
||||
build/
|
||||
dist/
|
||||
*.egg-info/
|
||||
.mypy_cache/
|
||||
|
||||
# Likely test artifact
|
||||
MagicMock/
|
||||
@@ -0,0 +1,62 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to this project will be documented in this file.
|
||||
|
||||
## [2.0.1] - 2026-10-20
|
||||
|
||||
### Fixed
|
||||
- `DiscordSendSaveImage` failing with `NameError: name 'build_metadata_section' is not defined` when sending metadata to Discord (#57)
|
||||
- `DiscordSendSaveVideo` failing to record Discord CDN URLs and crashing when `github_cdn_update` was enabled
|
||||
|
||||
## [2.0.0] - 2026-01-20
|
||||
|
||||
### Major Refactoring Release
|
||||
|
||||
Complete architectural refactoring to improve code organization, reduce duplication, and add bot features.
|
||||
|
||||
### Added
|
||||
- **Directory Structure**: New `nodes/`, `shared/`, `bot/` organization
|
||||
- **BaseDiscordNode**: Shared base class for image and video nodes (343 lines of reusable code)
|
||||
- **Shared Utilities**: 17 modular utility files in `shared/` package
|
||||
- `shared/discord/` - webhook client, message builder, CDN extractor
|
||||
- `shared/media/` - video encoder, format utils, image processing
|
||||
- `shared/workflow/` - sanitizer, prompt extractor, workflow builder
|
||||
- **Bot Features**:
|
||||
- WebSocket reconnection with exponential backoff
|
||||
- Error delivery to Discord users
|
||||
- `/template` commands (save, load, list, delete)
|
||||
- `/history` and `/rerun` commands
|
||||
- Config templates: `config.yaml.example`, `.env.example`
|
||||
|
||||
### Changed
|
||||
- **Code Reduction**: Total node code reduced by 620 lines (24%)
|
||||
- Image node: 986 → 836 lines
|
||||
- Video node: 1562 → 1092 lines
|
||||
- **Imports**: All imports now use `shared/` package instead of `discordsend_utils/`
|
||||
- **Video Encoding**: Extracted to `FFmpegEncoder` and `PILEncoder` classes
|
||||
|
||||
### Fixed
|
||||
- Bot startup bugs (BotConfig import, missing json import, permission imports)
|
||||
- SDXL workflow prompt extraction support
|
||||
- CDN URL redundant sends on 204 responses
|
||||
- Message builder metadata section formatting
|
||||
|
||||
### Removed
|
||||
- `discordsend_utils/` directory (replaced by `shared/`)
|
||||
- Obsolete documentation files
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Changed
|
||||
- **Performance**: Optimized image processing in `DiscordSendSaveImage` node.
|
||||
- Replaced redundant PIL-to-Numpy conversions with optimized Torch operations (~70% faster tensor processing).
|
||||
- Optimized JPEG encoding to use PIL directly, bypassing OpenCV conversion (~30% faster).
|
||||
|
||||
## [1.1.0] - 2026-01-10
|
||||
|
||||
### Added
|
||||
- **Image Node**: `DiscordSendSaveImage` for saving and sending images with advanced formatting.
|
||||
- **Video Node**: `DiscordSendSaveVideo` for converting image sequences to video (GIF, MP4, WebM, ProRes).
|
||||
- **GitHub Integration**: Automatically archive Discord CDN URLs to a GitHub repository.
|
||||
- **Testing**: Added initial test suite.
|
||||
- **Logging**: Implemented structured logging for better debugging.
|
||||
@@ -0,0 +1,62 @@
|
||||
# Contributing to ComfyUI-DiscordSend
|
||||
|
||||
First off, thank you for considering contributing to ComfyUI-DiscordSend! It's people like you that make this such a great tool.
|
||||
|
||||
## Code of Conduct
|
||||
|
||||
This project and everyone participating in it is governed by the [Contributor Covenant Code of Conduct](https://www.contributor-covenant.org/version/2/1/code_of_conduct/). By participating, you are expected to uphold this code.
|
||||
|
||||
## How Can I Contribute?
|
||||
|
||||
### Reporting Bugs
|
||||
|
||||
If you've noticed a bug, [open a new issue](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/issues/new)! It's generally best if you get confirmation of your bug this way before starting to code.
|
||||
|
||||
### Suggesting Enhancements
|
||||
|
||||
If you have a feature request, [open a new issue](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/issues/new)! It's generally best to get approval for your feature request this way before starting to code.
|
||||
|
||||
### Your First Code Contribution
|
||||
|
||||
Unsure where to begin contributing? You can start by looking through `good first issue` and `help wanted` issues:
|
||||
* [Good first issues](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/labels/good%20first%20issue) - issues which should only require a few lines of code, and a test or two.
|
||||
* [Help wanted issues](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/labels/help%20wanted) - issues which should be a bit more involved than `good first issue` issues.
|
||||
|
||||
### Pull Requests
|
||||
|
||||
1. **Fork & create a branch.** [Fork the repository](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/fork) and create a branch with a descriptive name. A good branch name would be (where issue #325 is the ticket you're working on): `git checkout -b 325-add-japanese-translations`
|
||||
2. **Set up your development environment.** To get started, you'll want to clone the repository and set up a virtual environment.
|
||||
```sh
|
||||
git clone https://github.com/AEmotionStudio/ComfyUI-DiscordSend.git
|
||||
cd ComfyUI-DiscordSend
|
||||
python -m venv venv
|
||||
source venv/bin/activate
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
3. **Implement your fix or feature.** At this point, you're ready to make your changes! Feel free to ask for help; everyone is a beginner at first.
|
||||
4. **Make a Pull Request.** At this point, you should switch back to your `master` branch and make sure it's up to date with the latest upstream version of the repository.
|
||||
```sh
|
||||
git remote add upstream git@github.com:AEmotionStudio/ComfyUI-DiscordSend.git
|
||||
git checkout master
|
||||
git pull upstream master
|
||||
```
|
||||
Then update your feature branch from your local copy of `master`, and push it!
|
||||
```sh
|
||||
git checkout 325-add-japanese-translations
|
||||
git rebase master
|
||||
git push --force-with-lease origin 325-add-japanese-translations
|
||||
```
|
||||
Finally, go to GitHub and [make a Pull Request](https://github.com/AEmotionStudio/ComfyUI-DiscordSend/compare).
|
||||
|
||||
## Styleguides
|
||||
|
||||
### Git Commit Messages
|
||||
|
||||
* Use the present tense ("Add feature" not "Added feature").
|
||||
* Use the imperative mood ("Move cursor to..." not "Moves cursor to...").
|
||||
* Limit the first line to 72 characters or less.
|
||||
* Reference issues and pull requests liberally after the first line.
|
||||
|
||||
### Python Styleguide
|
||||
|
||||
All Python code should adhere to [PEP 8](https://www.python.org/dev/peps/pep-0008/).
|
||||
+674
@@ -0,0 +1,674 @@
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 3, 29 June 2007
|
||||
|
||||
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
|
||||
Everyone is permitted to copy and distribute verbatim copies
|
||||
of this license document, but changing it is not allowed.
|
||||
|
||||
Preamble
|
||||
|
||||
The GNU General Public License is a free, copyleft license for
|
||||
software and other kinds of works.
|
||||
|
||||
The licenses for most software and other practical works are designed
|
||||
to take away your freedom to share and change the works. By contrast,
|
||||
the GNU General Public License is intended to guarantee your freedom to
|
||||
share and change all versions of a program--to make sure it remains free
|
||||
software for all its users. We, the Free Software Foundation, use the
|
||||
GNU General Public License for most of our software; it applies also to
|
||||
any other work released this way by its authors. You can apply it to
|
||||
your programs, too.
|
||||
|
||||
When we speak of free software, we are referring to freedom, not
|
||||
price. Our General Public Licenses are designed to make sure that you
|
||||
have the freedom to distribute copies of free software (and charge for
|
||||
them if you wish), that you receive source code or can get it if you
|
||||
want it, that you can change the software or use pieces of it in new
|
||||
free programs, and that you know you can do these things.
|
||||
|
||||
To protect your rights, we need to prevent others from denying you
|
||||
these rights or asking you to surrender the rights. Therefore, you have
|
||||
certain responsibilities if you distribute copies of the software, or if
|
||||
you modify it: responsibilities to respect the freedom of others.
|
||||
|
||||
For example, if you distribute copies of such a program, whether
|
||||
gratis or for a fee, you must pass on to the recipients the same
|
||||
freedoms that you received. You must make sure that they, too, receive
|
||||
or can get the source code. And you must show them these terms so they
|
||||
know their rights.
|
||||
|
||||
Developers that use the GNU GPL protect your rights with two steps:
|
||||
(1) assert copyright on the software, and (2) offer you this License
|
||||
giving you legal permission to copy, distribute and/or modify it.
|
||||
|
||||
For the developers' and authors' protection, the GPL clearly explains
|
||||
that there is no warranty for this free software. For both users' and
|
||||
authors' sake, the GPL requires that modified versions be marked as
|
||||
changed, so that their problems will not be attributed erroneously to
|
||||
authors of previous versions.
|
||||
|
||||
Some devices are designed to deny users access to install or run
|
||||
modified versions of the software inside them, although the manufacturer
|
||||
can do so. This is fundamentally incompatible with the aim of
|
||||
protecting users' freedom to change the software. The systematic
|
||||
pattern of such abuse occurs in the area of products for individuals to
|
||||
use, which is precisely where it is most unacceptable. Therefore, we
|
||||
have designed this version of the GPL to prohibit the practice for those
|
||||
products. If such problems arise substantially in other domains, we
|
||||
stand ready to extend this provision to those domains in future versions
|
||||
of the GPL, as needed to protect the freedom of users.
|
||||
|
||||
Finally, every program is threatened constantly by software patents.
|
||||
States should not allow patents to restrict development and use of
|
||||
software on general-purpose computers, but in those that do, we wish to
|
||||
avoid the special danger that patents applied to a free program could
|
||||
make it effectively proprietary. To prevent this, the GPL assures that
|
||||
patents cannot be used to render the program non-free.
|
||||
|
||||
The precise terms and conditions for copying, distribution and
|
||||
modification follow.
|
||||
|
||||
TERMS AND CONDITIONS
|
||||
|
||||
0. Definitions.
|
||||
|
||||
"This License" refers to version 3 of the GNU General Public License.
|
||||
|
||||
"Copyright" also means copyright-like laws that apply to other kinds of
|
||||
works, such as semiconductor masks.
|
||||
|
||||
"The Program" refers to any copyrightable work licensed under this
|
||||
License. Each licensee is addressed as "you". "Licensees" and
|
||||
"recipients" may be individuals or organizations.
|
||||
|
||||
To "modify" a work means to copy from or adapt all or part of the work
|
||||
in a fashion requiring copyright permission, other than the making of an
|
||||
exact copy. The resulting work is called a "modified version" of the
|
||||
earlier work or a work "based on" the earlier work.
|
||||
|
||||
A "covered work" means either the unmodified Program or a work based
|
||||
on the Program.
|
||||
|
||||
To "propagate" a work means to do anything with it that, without
|
||||
permission, would make you directly or secondarily liable for
|
||||
infringement under applicable copyright law, except executing it on a
|
||||
computer or modifying a private copy. Propagation includes copying,
|
||||
distribution (with or without modification), making available to the
|
||||
public, and in some countries other activities as well.
|
||||
|
||||
To "convey" a work means any kind of propagation that enables other
|
||||
parties to make or receive copies. Mere interaction with a user through
|
||||
a computer network, with no transfer of a copy, is not conveying.
|
||||
|
||||
An interactive user interface displays "Appropriate Legal Notices"
|
||||
to the extent that it includes a convenient and prominently visible
|
||||
feature that (1) displays an appropriate copyright notice, and (2)
|
||||
tells the user that there is no warranty for the work (except to the
|
||||
extent that warranties are provided), that licensees may convey the
|
||||
work under this License, and how to view a copy of this License. If
|
||||
the interface presents a list of user commands or options, such as a
|
||||
menu, a prominent item in the list meets this criterion.
|
||||
|
||||
1. Source Code.
|
||||
|
||||
The "source code" for a work means the preferred form of the work
|
||||
for making modifications to it. "Object code" means any non-source
|
||||
form of a work.
|
||||
|
||||
A "Standard Interface" means an interface that either is an official
|
||||
standard defined by a recognized standards body, or, in the case of
|
||||
interfaces specified for a particular programming language, one that
|
||||
is widely used among developers working in that language.
|
||||
|
||||
The "System Libraries" of an executable work include anything, other
|
||||
than the work as a whole, that (a) is included in the normal form of
|
||||
packaging a Major Component, but which is not part of that Major
|
||||
Component, and (b) serves only to enable use of the work with that
|
||||
Major Component, or to implement a Standard Interface for which an
|
||||
implementation is available to the public in source code form. A
|
||||
"Major Component", in this context, means a major essential component
|
||||
(kernel, window system, and so on) of the specific operating system
|
||||
(if any) on which the executable work runs, or a compiler used to
|
||||
produce the work, or an object code interpreter used to run it.
|
||||
|
||||
The "Corresponding Source" for a work in object code form means all
|
||||
the source code needed to generate, install, and (for an executable
|
||||
work) run the object code and to modify the work, including scripts to
|
||||
control those activities. However, it does not include the work's
|
||||
System Libraries, or general-purpose tools or generally available free
|
||||
programs which are used unmodified in performing those activities but
|
||||
which are not part of the work. For example, Corresponding Source
|
||||
includes interface definition files associated with source files for
|
||||
the work, and the source code for shared libraries and dynamically
|
||||
linked subprograms that the work is specifically designed to require,
|
||||
such as by intimate data communication or control flow between those
|
||||
subprograms and other parts of the work.
|
||||
|
||||
The Corresponding Source need not include anything that users
|
||||
can regenerate automatically from other parts of the Corresponding
|
||||
Source.
|
||||
|
||||
The Corresponding Source for a work in source code form is that
|
||||
same work.
|
||||
|
||||
2. Basic Permissions.
|
||||
|
||||
All rights granted under this License are granted for the term of
|
||||
copyright on the Program, and are irrevocable provided the stated
|
||||
conditions are met. This License explicitly affirms your unlimited
|
||||
permission to run the unmodified Program. The output from running a
|
||||
covered work is covered by this License only if the output, given its
|
||||
content, constitutes a covered work. This License acknowledges your
|
||||
rights of fair use or other equivalent, as provided by copyright law.
|
||||
|
||||
You may make, run and propagate covered works that you do not
|
||||
convey, without conditions so long as your license otherwise remains
|
||||
in force. You may convey covered works to others for the sole purpose
|
||||
of having them make modifications exclusively for you, or provide you
|
||||
with facilities for running those works, provided that you comply with
|
||||
the terms of this License in conveying all material for which you do
|
||||
not control copyright. Those thus making or running the covered works
|
||||
for you must do so exclusively on your behalf, under your direction
|
||||
and control, on terms that prohibit them from making any copies of
|
||||
your copyrighted material outside their relationship with you.
|
||||
|
||||
Conveying under any other circumstances is permitted solely under
|
||||
the conditions stated below. Sublicensing is not allowed; section 10
|
||||
makes it unnecessary.
|
||||
|
||||
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
|
||||
|
||||
No covered work shall be deemed part of an effective technological
|
||||
measure under any applicable law fulfilling obligations under article
|
||||
11 of the WIPO copyright treaty adopted on 20 December 1996, or
|
||||
similar laws prohibiting or restricting circumvention of such
|
||||
measures.
|
||||
|
||||
When you convey a covered work, you waive any legal power to forbid
|
||||
circumvention of technological measures to the extent such circumvention
|
||||
is effected by exercising rights under this License with respect to
|
||||
the covered work, and you disclaim any intention to limit operation or
|
||||
modification of the work as a means of enforcing, against the work's
|
||||
users, your or third parties' legal rights to forbid circumvention of
|
||||
technological measures.
|
||||
|
||||
4. Conveying Verbatim Copies.
|
||||
|
||||
You may convey verbatim copies of the Program's source code as you
|
||||
receive it, in any medium, provided that you conspicuously and
|
||||
appropriately publish on each copy an appropriate copyright notice;
|
||||
keep intact all notices stating that this License and any
|
||||
non-permissive terms added in accord with section 7 apply to the code;
|
||||
keep intact all notices of the absence of any warranty; and give all
|
||||
recipients a copy of this License along with the Program.
|
||||
|
||||
You may charge any price or no price for each copy that you convey,
|
||||
and you may offer support or warranty protection for a fee.
|
||||
|
||||
5. Conveying Modified Source Versions.
|
||||
|
||||
You may convey a work based on the Program, or the modifications to
|
||||
produce it from the Program, in the form of source code under the
|
||||
terms of section 4, provided that you also meet all of these conditions:
|
||||
|
||||
a) The work must carry prominent notices stating that you modified
|
||||
it, and giving a relevant date.
|
||||
|
||||
b) The work must carry prominent notices stating that it is
|
||||
released under this License and any conditions added under section
|
||||
7. This requirement modifies the requirement in section 4 to
|
||||
"keep intact all notices".
|
||||
|
||||
c) You must license the entire work, as a whole, under this
|
||||
License to anyone who comes into possession of a copy. This
|
||||
License will therefore apply, along with any applicable section 7
|
||||
additional terms, to the whole of the work, and all its parts,
|
||||
regardless of how they are packaged. This License gives no
|
||||
permission to license the work in any other way, but it does not
|
||||
invalidate such permission if you have separately received it.
|
||||
|
||||
d) If the work has interactive user interfaces, each must display
|
||||
Appropriate Legal Notices; however, if the Program has interactive
|
||||
interfaces that do not display Appropriate Legal Notices, your
|
||||
work need not make them do so.
|
||||
|
||||
A compilation of a covered work with other separate and independent
|
||||
works, which are not by their nature extensions of the covered work,
|
||||
and which are not combined with it such as to form a larger program,
|
||||
in or on a volume of a storage or distribution medium, is called an
|
||||
"aggregate" if the compilation and its resulting copyright are not
|
||||
used to limit the access or legal rights of the compilation's users
|
||||
beyond what the individual works permit. Inclusion of a covered work
|
||||
in an aggregate does not cause this License to apply to the other
|
||||
parts of the aggregate.
|
||||
|
||||
6. Conveying Non-Source Forms.
|
||||
|
||||
You may convey a covered work in object code form under the terms
|
||||
of sections 4 and 5, provided that you also convey the
|
||||
machine-readable Corresponding Source under the terms of this License,
|
||||
in one of these ways:
|
||||
|
||||
a) Convey the object code in, or embodied in, a physical product
|
||||
(including a physical distribution medium), accompanied by the
|
||||
Corresponding Source fixed on a durable physical medium
|
||||
customarily used for software interchange.
|
||||
|
||||
b) Convey the object code in, or embodied in, a physical product
|
||||
(including a physical distribution medium), accompanied by a
|
||||
written offer, valid for at least three years and valid for as
|
||||
long as you offer spare parts or customer support for that product
|
||||
model, to give anyone who possesses the object code either (1) a
|
||||
copy of the Corresponding Source for all the software in the
|
||||
product that is covered by this License, on a durable physical
|
||||
medium customarily used for software interchange, for a price no
|
||||
more than your reasonable cost of physically performing this
|
||||
conveying of source, or (2) access to copy the
|
||||
Corresponding Source from a network server at no charge.
|
||||
|
||||
c) Convey individual copies of the object code with a copy of the
|
||||
written offer to provide the Corresponding Source. This
|
||||
alternative is allowed only occasionally and noncommercially, and
|
||||
only if you received the object code with such an offer, in accord
|
||||
with subsection 6b.
|
||||
|
||||
d) Convey the object code by offering access from a designated
|
||||
place (gratis or for a charge), and offer equivalent access to the
|
||||
Corresponding Source in the same way through the same place at no
|
||||
further charge. You need not require recipients to copy the
|
||||
Corresponding Source along with the object code. If the place to
|
||||
copy the object code is a network server, the Corresponding Source
|
||||
may be on a different server (operated by you or a third party)
|
||||
that supports equivalent copying facilities, provided you maintain
|
||||
clear directions next to the object code saying where to find the
|
||||
Corresponding Source. Regardless of what server hosts the
|
||||
Corresponding Source, you remain obligated to ensure that it is
|
||||
available for as long as needed to satisfy these requirements.
|
||||
|
||||
e) Convey the object code using peer-to-peer transmission, provided
|
||||
you inform other peers where the object code and Corresponding
|
||||
Source of the work are being offered to the general public at no
|
||||
charge under subsection 6d.
|
||||
|
||||
A separable portion of the object code, whose source code is excluded
|
||||
from the Corresponding Source as a System Library, need not be
|
||||
included in conveying the object code work.
|
||||
|
||||
A "User Product" is either (1) a "consumer product", which means any
|
||||
tangible personal property which is normally used for personal, family,
|
||||
or household purposes, or (2) anything designed or sold for incorporation
|
||||
into a dwelling. In determining whether a product is a consumer product,
|
||||
doubtful cases shall be resolved in favor of coverage. For a particular
|
||||
product received by a particular user, "normally used" refers to a
|
||||
typical or common use of that class of product, regardless of the status
|
||||
of the particular user or of the way in which the particular user
|
||||
actually uses, or expects or is expected to use, the product. A product
|
||||
is a consumer product regardless of whether the product has substantial
|
||||
commercial, industrial or non-consumer uses, unless such uses represent
|
||||
the only significant mode of use of the product.
|
||||
|
||||
"Installation Information" for a User Product means any methods,
|
||||
procedures, authorization keys, or other information required to install
|
||||
and execute modified versions of a covered work in that User Product from
|
||||
a modified version of its Corresponding Source. The information must
|
||||
suffice to ensure that the continued functioning of the modified object
|
||||
code is in no case prevented or interfered with solely because
|
||||
modification has been made.
|
||||
|
||||
If you convey an object code work under this section in, or with, or
|
||||
specifically for use in, a User Product, and the conveying occurs as
|
||||
part of a transaction in which the right of possession and use of the
|
||||
User Product is transferred to the recipient in perpetuity or for a
|
||||
fixed term (regardless of how the transaction is characterized), the
|
||||
Corresponding Source conveyed under this section must be accompanied
|
||||
by the Installation Information. But this requirement does not apply
|
||||
if neither you nor any third party retains the ability to install
|
||||
modified object code on the User Product (for example, the work has
|
||||
been installed in ROM).
|
||||
|
||||
The requirement to provide Installation Information does not include a
|
||||
requirement to continue to provide support service, warranty, or updates
|
||||
for a work that has been modified or installed by the recipient, or for
|
||||
the User Product in which it has been modified or installed. Access to a
|
||||
network may be denied when the modification itself materially and
|
||||
adversely affects the operation of the network or violates the rules and
|
||||
protocols for communication across the network.
|
||||
|
||||
Corresponding Source conveyed, and Installation Information provided,
|
||||
in accord with this section must be in a format that is publicly
|
||||
documented (and with an implementation available to the public in
|
||||
source code form), and must require no special password or key for
|
||||
unpacking, reading or copying.
|
||||
|
||||
7. Additional Terms.
|
||||
|
||||
"Additional permissions" are terms that supplement the terms of this
|
||||
License by making exceptions from one or more of its conditions.
|
||||
Additional permissions that are applicable to the entire Program shall
|
||||
be treated as though they were included in this License, to the extent
|
||||
that they are valid under applicable law. If additional permissions
|
||||
apply only to part of the Program, that part may be used separately
|
||||
under those permissions, but the entire Program remains governed by
|
||||
this License without regard to the additional permissions.
|
||||
|
||||
When you convey a copy of a covered work, you may at your option
|
||||
remove any additional permissions from that copy, or from any part of
|
||||
it. (Additional permissions may be written to require their own
|
||||
removal in certain cases when you modify the work.) You may place
|
||||
additional permissions on material, added by you to a covered work,
|
||||
for which you have or can give appropriate copyright permission.
|
||||
|
||||
Notwithstanding any other provision of this License, for material you
|
||||
add to a covered work, you may (if authorized by the copyright holders of
|
||||
that material) supplement the terms of this License with terms:
|
||||
|
||||
a) Disclaiming warranty or limiting liability differently from the
|
||||
terms of sections 15 and 16 of this License; or
|
||||
|
||||
b) Requiring preservation of specified reasonable legal notices or
|
||||
author attributions in that material or in the Appropriate Legal
|
||||
Notices displayed by works containing it; or
|
||||
|
||||
c) Prohibiting misrepresentation of the origin of that material, or
|
||||
requiring that modified versions of such material be marked in
|
||||
reasonable ways as different from the original version; or
|
||||
|
||||
d) Limiting the use for publicity purposes of names of licensors or
|
||||
authors of the material; or
|
||||
|
||||
e) Declining to grant rights under trademark law for use of some
|
||||
trade names, trademarks, or service marks; or
|
||||
|
||||
f) Requiring indemnification of licensors and authors of that
|
||||
material by anyone who conveys the material (or modified versions of
|
||||
it) with contractual assumptions of liability to the recipient, for
|
||||
any liability that these contractual assumptions directly impose on
|
||||
those licensors and authors.
|
||||
|
||||
All other non-permissive additional terms are considered "further
|
||||
restrictions" within the meaning of section 10. If the Program as you
|
||||
received it, or any part of it, contains a notice stating that it is
|
||||
governed by this License along with a term that is a further
|
||||
restriction, you may remove that term. If a license document contains
|
||||
a further restriction but permits relicensing or conveying under this
|
||||
License, you may add to a covered work material governed by the terms
|
||||
of that license document, provided that the further restriction does
|
||||
not survive such relicensing or conveying.
|
||||
|
||||
If you add terms to a covered work in accord with this section, you
|
||||
must place, in the relevant source files, a statement of the
|
||||
additional terms that apply to those files, or a notice indicating
|
||||
where to find the applicable terms.
|
||||
|
||||
Additional terms, permissive or non-permissive, may be stated in the
|
||||
form of a separately written license, or stated as exceptions;
|
||||
the above requirements apply either way.
|
||||
|
||||
8. Termination.
|
||||
|
||||
You may not propagate or modify a covered work except as expressly
|
||||
provided under this License. Any attempt otherwise to propagate or
|
||||
modify it is void, and will automatically terminate your rights under
|
||||
this License (including any patent licenses granted under the third
|
||||
paragraph of section 11).
|
||||
|
||||
However, if you cease all violation of this License, then your
|
||||
license from a particular copyright holder is reinstated (a)
|
||||
provisionally, unless and until the copyright holder explicitly and
|
||||
finally terminates your license, and (b) permanently, if the copyright
|
||||
holder fails to notify you of the violation by some reasonable means
|
||||
prior to 60 days after the cessation.
|
||||
|
||||
Moreover, your license from a particular copyright holder is
|
||||
reinstated permanently if the copyright holder notifies you of the
|
||||
violation by some reasonable means, this is the first time you have
|
||||
received notice of violation of this License (for any work) from that
|
||||
copyright holder, and you cure the violation prior to 30 days after
|
||||
your receipt of the notice.
|
||||
|
||||
Termination of your rights under this section does not terminate the
|
||||
licenses of parties who have received copies or rights from you under
|
||||
this License. If your rights have been terminated and not permanently
|
||||
reinstated, you do not qualify to receive new licenses for the same
|
||||
material under section 10.
|
||||
|
||||
9. Acceptance Not Required for Having Copies.
|
||||
|
||||
You are not required to accept this License in order to receive or
|
||||
run a copy of the Program. Ancillary propagation of a covered work
|
||||
occurring solely as a consequence of using peer-to-peer transmission
|
||||
to receive a copy likewise does not require acceptance. However,
|
||||
nothing other than this License grants you permission to propagate or
|
||||
modify any covered work. These actions infringe copyright if you do
|
||||
not accept this License. Therefore, by modifying or propagating a
|
||||
covered work, you indicate your acceptance of this License to do so.
|
||||
|
||||
10. Automatic Licensing of Downstream Recipients.
|
||||
|
||||
Each time you convey a covered work, the recipient automatically
|
||||
receives a license from the original licensors, to run, modify and
|
||||
propagate that work, subject to this License. You are not responsible
|
||||
for enforcing compliance by third parties with this License.
|
||||
|
||||
An "entity transaction" is a transaction transferring control of an
|
||||
organization, or substantially all assets of one, or subdividing an
|
||||
organization, or merging organizations. If propagation of a covered
|
||||
work results from an entity transaction, each party to that
|
||||
transaction who receives a copy of the work also receives whatever
|
||||
licenses to the work the party's predecessor in interest had or could
|
||||
give under the previous paragraph, plus a right to possession of the
|
||||
Corresponding Source of the work from the predecessor in interest, if
|
||||
the predecessor has it or can get it with reasonable efforts.
|
||||
|
||||
You may not impose any further restrictions on the exercise of the
|
||||
rights granted or affirmed under this License. For example, you may
|
||||
not impose a license fee, royalty, or other charge for exercise of
|
||||
rights granted under this License, and you may not initiate litigation
|
||||
(including a cross-claim or counterclaim in a lawsuit) alleging that
|
||||
any patent claim is infringed by making, using, selling, offering for
|
||||
sale, or importing the Program or any portion of it.
|
||||
|
||||
11. Patents.
|
||||
|
||||
A "contributor" is a copyright holder who authorizes use under this
|
||||
License of the Program or a work on which the Program is based. The
|
||||
work thus licensed is called the contributor's "contributor version".
|
||||
|
||||
A contributor's "essential patent claims" are all patent claims
|
||||
owned or controlled by the contributor, whether already acquired or
|
||||
hereafter acquired, that would be infringed by some manner, permitted
|
||||
by this License, of making, using, or selling its contributor version,
|
||||
but do not include claims that would be infringed only as a
|
||||
consequence of further modification of the contributor version. For
|
||||
purposes of this definition, "control" includes the right to grant
|
||||
patent sublicenses in a manner consistent with the requirements of
|
||||
this License.
|
||||
|
||||
Each contributor grants you a non-exclusive, worldwide, royalty-free
|
||||
patent license under the contributor's essential patent claims, to
|
||||
make, use, sell, offer for sale, import and otherwise run, modify and
|
||||
propagate the contents of its contributor version.
|
||||
|
||||
In the following three paragraphs, a "patent license" is any express
|
||||
agreement or commitment, however denominated, not to enforce a patent
|
||||
(such as an express permission to practice a patent or covenant not to
|
||||
sue for patent infringement). To "grant" such a patent license to a
|
||||
party means to make such an agreement or commitment not to enforce a
|
||||
patent against the party.
|
||||
|
||||
If you convey a covered work, knowingly relying on a patent license,
|
||||
and the Corresponding Source of the work is not available for anyone
|
||||
to copy, free of charge and under the terms of this License, through a
|
||||
publicly available network server or other readily accessible means,
|
||||
then you must either (1) cause the Corresponding Source to be so
|
||||
available, or (2) arrange to deprive yourself of the benefit of the
|
||||
patent license for this particular work, or (3) arrange, in a manner
|
||||
consistent with the requirements of this License, to extend the patent
|
||||
license to downstream recipients. "Knowingly relying" means you have
|
||||
actual knowledge that, but for the patent license, your conveying the
|
||||
covered work in a country, or your recipient's use of the covered work
|
||||
in a country, would infringe one or more identifiable patents in that
|
||||
country that you have reason to believe are valid.
|
||||
|
||||
If, pursuant to or in connection with a single transaction or
|
||||
arrangement, you convey, or propagate by procuring conveyance of, a
|
||||
covered work, and grant a patent license to some of the parties
|
||||
receiving the covered work authorizing them to use, propagate, modify
|
||||
or convey a specific copy of the covered work, then the patent license
|
||||
you grant is automatically extended to all recipients of the covered
|
||||
work and works based on it.
|
||||
|
||||
A patent license is "discriminatory" if it does not include within
|
||||
the scope of its coverage, prohibits the exercise of, or is
|
||||
conditioned on the non-exercise of one or more of the rights that are
|
||||
specifically granted under this License. You may not convey a covered
|
||||
work if you are a party to an arrangement with a third party that is
|
||||
in the business of distributing software, under which you make payment
|
||||
to the third party based on the extent of your activity of conveying
|
||||
the work, and under which the third party grants, to any of the
|
||||
parties who would receive the covered work from you, a discriminatory
|
||||
patent license (a) in connection with copies of the covered work
|
||||
conveyed by you (or copies made from those copies), or (b) primarily
|
||||
for and in connection with specific products or compilations that
|
||||
contain the covered work, unless you entered into that arrangement,
|
||||
or that patent license was granted, prior to 28 March 2007.
|
||||
|
||||
Nothing in this License shall be construed as excluding or limiting
|
||||
any implied license or other defenses to infringement that may
|
||||
otherwise be available to you under applicable patent law.
|
||||
|
||||
12. No Surrender of Others' Freedom.
|
||||
|
||||
If conditions are imposed on you (whether by court order, agreement or
|
||||
otherwise) that contradict the conditions of this License, they do not
|
||||
excuse you from the conditions of this License. If you cannot convey a
|
||||
covered work so as to satisfy simultaneously your obligations under this
|
||||
License and any other pertinent obligations, then as a consequence you may
|
||||
not convey it at all. For example, if you agree to terms that obligate you
|
||||
to collect a royalty for further conveying from those to whom you convey
|
||||
the Program, the only way you could satisfy both those terms and this
|
||||
License would be to refrain entirely from conveying the Program.
|
||||
|
||||
13. Use with the GNU Affero General Public License.
|
||||
|
||||
Notwithstanding any other provision of this License, you have
|
||||
permission to link or combine any covered work with a work licensed
|
||||
under version 3 of the GNU Affero General Public License into a single
|
||||
combined work, and to convey the resulting work. The terms of this
|
||||
License will continue to apply to the part which is the covered work,
|
||||
but the special requirements of the GNU Affero General Public License,
|
||||
section 13, concerning interaction through a network will apply to the
|
||||
combination as such.
|
||||
|
||||
14. Revised Versions of this License.
|
||||
|
||||
The Free Software Foundation may publish revised and/or new versions of
|
||||
the GNU General Public License from time to time. Such new versions will
|
||||
be similar in spirit to the present version, but may differ in detail to
|
||||
address new problems or concerns.
|
||||
|
||||
Each version is given a distinguishing version number. If the
|
||||
Program specifies that a certain numbered version of the GNU General
|
||||
Public License "or any later version" applies to it, you have the
|
||||
option of following the terms and conditions either of that numbered
|
||||
version or of any later version published by the Free Software
|
||||
Foundation. If the Program does not specify a version number of the
|
||||
GNU General Public License, you may choose any version ever published
|
||||
by the Free Software Foundation.
|
||||
|
||||
If the Program specifies that a proxy can decide which future
|
||||
versions of the GNU General Public License can be used, that proxy's
|
||||
public statement of acceptance of a version permanently authorizes you
|
||||
to choose that version for the Program.
|
||||
|
||||
Later license versions may give you additional or different
|
||||
permissions. However, no additional obligations are imposed on any
|
||||
author or copyright holder as a result of your choosing to follow a
|
||||
later version.
|
||||
|
||||
15. Disclaimer of Warranty.
|
||||
|
||||
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
|
||||
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
|
||||
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
|
||||
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
|
||||
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
|
||||
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
|
||||
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
|
||||
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
|
||||
|
||||
16. Limitation of Liability.
|
||||
|
||||
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
|
||||
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
|
||||
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
|
||||
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
|
||||
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
|
||||
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
|
||||
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
|
||||
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
|
||||
SUCH DAMAGES.
|
||||
|
||||
17. Interpretation of Sections 15 and 16.
|
||||
|
||||
If the disclaimer of warranty and limitation of liability provided
|
||||
above cannot be given local legal effect according to their terms,
|
||||
reviewing courts shall apply local law that most closely approximates
|
||||
an absolute waiver of all civil liability in connection with the
|
||||
Program, unless a warranty or assumption of liability accompanies a
|
||||
copy of the Program in return for a fee.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
How to Apply These Terms to Your New Programs
|
||||
|
||||
If you develop a new program, and you want it to be of the greatest
|
||||
possible use to the public, the best way to achieve this is to make it
|
||||
free software which everyone can redistribute and change under these terms.
|
||||
|
||||
To do so, attach the following notices to the program. It is safest
|
||||
to attach them to the start of each source file to most effectively
|
||||
state the exclusion of warranty; and each file should have at least
|
||||
the "copyright" line and a pointer to where the full notice is found.
|
||||
|
||||
<one line to give the program's name and a brief idea of what it does.>
|
||||
Copyright (C) <year> <name of author>
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
Also add information on how to contact you by electronic and paper mail.
|
||||
|
||||
If the program does terminal interaction, make it output a short
|
||||
notice like this when it starts in an interactive mode:
|
||||
|
||||
<program> Copyright (C) <year> <name of author>
|
||||
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
|
||||
This is free software, and you are welcome to redistribute it
|
||||
under certain conditions; type `show c' for details.
|
||||
|
||||
The hypothetical commands `show w' and `show c' should show the appropriate
|
||||
parts of the General Public License. Of course, your program's commands
|
||||
might be different; for a GUI interface, you would use an "about box".
|
||||
|
||||
You should also get your employer (if you work as a programmer) or school,
|
||||
if any, to sign a "copyright disclaimer" for the program, if necessary.
|
||||
For more information on this, and how to apply and follow the GNU GPL, see
|
||||
<https://www.gnu.org/licenses/>.
|
||||
|
||||
The GNU General Public License does not permit incorporating your program
|
||||
into proprietary programs. If your program is a subroutine library, you
|
||||
may consider it more useful to permit linking proprietary applications with
|
||||
the library. If this is what you want to do, use the GNU Lesser General
|
||||
Public License instead of this License. But first, please read
|
||||
<https://www.gnu.org/licenses/why-not-lgpl.html>.
|
||||
@@ -0,0 +1,323 @@
|
||||
<div align="center">
|
||||
|
||||
# ComfyUI-DiscordSend
|
||||
|
||||

|
||||

|
||||

|
||||

|
||||
<br>
|
||||

|
||||

|
||||

|
||||
<br>
|
||||

|
||||

|
||||
|
||||

|
||||
|
||||
**Send your AI-generated images and videos directly to Discord from ComfyUI!**
|
||||
|
||||
[Features](#-features) • [Installation](#-installation) • [Settings](#-settings) • [Technical Details](#-technical-details) • [Contributing](#-contributing) • [Changelog](#-changelog)
|
||||
|
||||
</div>
|
||||
|
||||
---
|
||||
|
||||
## What's New in v2.0.0 (January 20, 2026)
|
||||
|
||||
### 🚀 Major Refactoring Release
|
||||
- **New Architecture**: Reorganized into `nodes/`, `shared/`, and `bot/` packages for cleaner separation of concerns.
|
||||
- **BaseDiscordNode**: Shared base class eliminates 620 lines of duplicated code across image and video nodes.
|
||||
- **Shared Utilities**: 17 modular utility files covering Discord webhooks, media processing, and workflow handling.
|
||||
- **Discord Bot** *(optional)*: New standalone bot with slash commands, WebSocket reconnection, and job queuing.
|
||||
- **Performance**: Optimized image processing with direct Torch operations (~70% faster tensor processing).
|
||||
|
||||
📄 See [CHANGELOG.md](CHANGELOG.md) for the complete version history.
|
||||
|
||||
---
|
||||
|
||||
## ✨ Features
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="60%">
|
||||
<h3>🖼️ Image Node: DiscordSendSaveImage</h3>
|
||||
<ul>
|
||||
<li>Save images in various formats (PNG, JPEG, WebP)</li>
|
||||
<li>Send images directly to Discord via webhooks</li>
|
||||
<li>Include workflow JSON for easy reproduction</li>
|
||||
<li>Customizable Discord messages with prompt information</li>
|
||||
<li>Advanced file naming with date, time, and dimension options</li>
|
||||
<li>High-quality image export with configurable compression settings</li>
|
||||
<li>Built-in preview functionality within ComfyUI interface</li>
|
||||
<li>Batch grouping support (up to 9 images per Discord message)</li>
|
||||
<li>Unique identifier support to prevent conflicts between users</li>
|
||||
</ul>
|
||||
</td>
|
||||
<td width="40%">
|
||||
<img src="/images/nodes_example_3.png" alt="Image Node" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="40%">
|
||||
<img src="/images/discord_formatting.png" alt="Video Node" width="100%">
|
||||
</td>
|
||||
<td width="60%">
|
||||
<h3>🎬 Video Node: DiscordSendSaveVideo</h3>
|
||||
<ul>
|
||||
<li>Convert image sequences to videos in multiple formats</li>
|
||||
<li>Support for GIF, MP4, WebM, and professional formats like ProRes</li>
|
||||
<li>Configurable frame rates from 0.1 to 120 fps</li>
|
||||
<li>Extra-slow frame rates for photo slideshows</li>
|
||||
<li>Add audio to your videos</li>
|
||||
<li>Special effects like ping-pong looping</li>
|
||||
<li>Discord integration for sharing videos</li>
|
||||
<li>Include workflow data and video information in messages</li>
|
||||
</ul>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="60%">
|
||||
<h3>🔄 GitHub Integration</h3>
|
||||
<ul>
|
||||
<li>Save Discord CDN URLs to a GitHub repository</li>
|
||||
<li>Automatically update existing URL collections with new uploads</li>
|
||||
<li>Formatted markdown files with timestamps and organized links</li>
|
||||
<li>Perfect for building media galleries or documentation</li>
|
||||
<li>Comprehensive security measures for GitHub tokens</li>
|
||||
</ul>
|
||||
</td>
|
||||
<td width="40%">
|
||||
<img src="/images/discord_webhook_step2.png" alt="GitHub Integration" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
## 📥 Installation
|
||||
|
||||
### Option 1: Using ComfyUI Manager
|
||||
|
||||
1. Install [ComfyUI Manager](https://github.com/ltdrdata/ComfyUI-Manager) if you don't have it already
|
||||
2. Open ComfyUI, go to the Manager tab
|
||||
3. Search for "ComfyUI-DiscordSend" and install
|
||||
|
||||
### Option 2: Manual Installation
|
||||
|
||||
```bash
|
||||
cd /path/to/ComfyUI/custom_nodes
|
||||
git clone https://github.com/AEmotionStudio/ComfyUI-DiscordSend
|
||||
cd ComfyUI-DiscordSend
|
||||
|
||||
# For nodes only (minimal - just the requests library):
|
||||
pip install -r requirements-nodes.txt
|
||||
|
||||
# For full bot support (Discord bot + all features):
|
||||
pip install -r requirements-bot.txt
|
||||
```
|
||||
|
||||
> [!IMPORTANT]
|
||||
> - For video functionality, ffmpeg must be installed on your system. The node will automatically detect its presence.
|
||||
> - **Nodes only** require just the `requests` library (1 dependency).
|
||||
> - **Discord bot** requires additional dependencies (discord.py, aiohttp, sqlalchemy, etc.).
|
||||
|
||||
## 🏗️ Architecture Overview
|
||||
|
||||
This extension contains **two independent systems**. Most users only need the nodes.
|
||||
|
||||
| | ComfyUI Nodes | Discord Bot |
|
||||
|---|---|---|
|
||||
| **What it does** | Send images/videos to Discord from your workflow | Standalone bot that lets Discord users trigger ComfyUI workflows via slash commands |
|
||||
| **How to use** | Add `DiscordSendSaveImage` or `DiscordSendSaveVideo` to your workflow | Run `python -m bot` as a separate process alongside ComfyUI |
|
||||
| **Auth method** | Paste a **webhook URL** directly into the node | Requires a **bot token** from the Discord Developer Portal |
|
||||
| **Config files needed** | ❌ None — everything is configured in the node itself | ✅ `.env` or `config.yaml` (copy from `.env.example` / `config.yaml.example`) |
|
||||
| **Dependencies** | `requests` only | `discord.py`, `aiohttp`, `sqlalchemy`, etc. |
|
||||
|
||||
> [!NOTE]
|
||||
> The `.env.example` and `config.yaml.example` files in the repository root are **only for the optional Discord bot**. If you just want to send images/videos to Discord from your ComfyUI workflows, you do not need these files. Simply paste your Discord webhook URL into the node's `webhook_url` field.
|
||||
|
||||
## ⚙️ Settings
|
||||
|
||||
### Discord Webhook Setup
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="50%">
|
||||
<h3>Step 1: Access Webhook Settings</h3>
|
||||
<ol>
|
||||
<li>In Discord, go to your server settings → Integrations → Webhooks</li>
|
||||
</ol>
|
||||
</td>
|
||||
<td width="50%">
|
||||
<img src="/images/discord_webhook_step1.png" alt="Discord webhook settings" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="50%">
|
||||
<h3>Step 2: Create and Copy Webhook</h3>
|
||||
<ol start="2">
|
||||
<li>Create a New Webhook and copy the Webhook URL</li>
|
||||
</ol>
|
||||
</td>
|
||||
<td width="50%">
|
||||
<img src="/images/discord_webhook_step2.png" alt="Create webhook" width="100%">
|
||||
<br/><br/>
|
||||
<img src="/images/discord_webhook_step2_2.png" alt="Copy webhook" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<table>
|
||||
<tr>
|
||||
<td width="50%">
|
||||
<h3>Step 3: Configure the Node</h3>
|
||||
<ol start="3">
|
||||
<li>Paste this URL into the "webhook_url" field in the DiscordSend nodes</li>
|
||||
</ol>
|
||||
</td>
|
||||
<td width="50%">
|
||||
<img src="/images/discord_webhook_step4.png" alt="Configure webhook in node" width="100%">
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
<details>
|
||||
<summary><strong>🖼️ Image Node Options (Click to Expand)</strong></summary>
|
||||
|
||||
| Option | Description |
|
||||
|--------|-------------|
|
||||
| **Required Parameters** ||
|
||||
| `images` | The images to save and/or send to Discord |
|
||||
| `filename_prefix` | Prefix for saved files (default: "ComfyUI-Image") |
|
||||
| `overwrite_last` | Enable to overwrite last image instead of incrementing filenames |
|
||||
| **File Options** ||
|
||||
| `file_format` | PNG, JPEG, or WebP |
|
||||
| `quality` | 1-100 for JPEG/WebP formats |
|
||||
| `lossless` | Use lossless compression when available |
|
||||
| `save_output` | Whether to save images to disk |
|
||||
| `show_preview` | Show preview in ComfyUI interface |
|
||||
| **Filename Options** ||
|
||||
| `add_date` | Add current date to filenames |
|
||||
| `add_time` | Add current time to filenames |
|
||||
| `add_dimensions` | Add width and height to filenames |
|
||||
| `resize_to_power_of_2` | Resize to nearest power of 2 dimensions |
|
||||
| `resize_method` | Method for resizing (nearest, bilinear, etc.) |
|
||||
| **Discord Options** ||
|
||||
| `send_to_discord` | Enable Discord webhook integration |
|
||||
| `webhook_url` | Discord webhook URL |
|
||||
| `discord_message` | Text message to accompany images |
|
||||
| `include_prompts_in_message` | Include generation prompts in Discord message |
|
||||
| `include_format_in_message` | Include image format details in message |
|
||||
| `group_batched_images` | Group batch images into one Discord message (max 9 images) |
|
||||
| `send_workflow_json` | Send workflow JSON for reproducibility |
|
||||
| **GitHub Options** ||
|
||||
| `save_cdn_urls` | Save the Discord CDN URLs as a text file and attach to Discord message |
|
||||
| `github_cdn_update` | Update a GitHub repository with Discord CDN URLs |
|
||||
| `github_repo` | GitHub repository in format "username/repo" |
|
||||
| `github_token` | GitHub personal access token (with repo permissions) |
|
||||
| `github_file_path` | Path to file in repository to update (default: "cdn_urls.md") |
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><strong>🎬 Video Node Options (Click to Expand)</strong></summary>
|
||||
|
||||
| Option | Description |
|
||||
|--------|-------------|
|
||||
| **Required Parameters** ||
|
||||
| `images` | Image sequence to convert to video |
|
||||
| `filename_prefix` | Prefix for saved files (default: "ComfyUI-Video") |
|
||||
| `overwrite_last` | Enable to overwrite last video instead of incrementing filenames |
|
||||
| **Video Options** ||
|
||||
| `format` | Various formats including GIF, MP4, WebM, ProRes |
|
||||
| `frame_rate` | Frames per second (0.1-120), values below 1 make images stay longer |
|
||||
| `quality` | Quality setting for compression (1-100) |
|
||||
| `loop_count` | Number of loops for GIF (0=infinite) |
|
||||
| `lossless` | Use lossless compression when available |
|
||||
| `pingpong` | Create forward-backward looping effect |
|
||||
| `save_output` | Whether to save video to disk |
|
||||
| `audio` | Optional audio to embed in video |
|
||||
| **Filename Options** ||
|
||||
| `add_date` | Add current date to filenames |
|
||||
| `add_time` | Add current time to filenames (DO NOT DISABLE for Discord uploads - see known issues) |
|
||||
| `add_dimensions` | Add width and height to filenames |
|
||||
| **Discord Options** ||
|
||||
| `send_to_discord` | Enable Discord webhook integration |
|
||||
| `webhook_url` | Discord webhook URL |
|
||||
| `discord_message` | Text message to accompany videos |
|
||||
| `include_prompts_in_message` | Include generation prompts in Discord message |
|
||||
| `include_video_info` | Include video format details in message (disable if you don't want time/format info) |
|
||||
| `send_workflow_json` | Send workflow JSON for reproducibility |
|
||||
| **GitHub Options** ||
|
||||
| `save_cdn_urls` | Save the Discord CDN URLs as a text file and attach to Discord message |
|
||||
| `github_cdn_update` | Update a GitHub repository with Discord CDN URLs |
|
||||
| `github_repo` | GitHub repository in format "username/repo" |
|
||||
| `github_token` | GitHub personal access token (with repo permissions) |
|
||||
| `github_file_path` | Path to file in repository to update (default: "cdn_urls.md") |
|
||||
|
||||
</details>
|
||||
|
||||
## 🔧 Technical Details
|
||||
|
||||
| Component | Description |
|
||||
|-----------|-------------|
|
||||
| **Video Processing** | Discord reprocesses all uploaded videos. Quality settings may be modified, and videos compressed. |
|
||||
| **Max Resolution** | Typically limited to 1080p by Discord. |
|
||||
| **Max Bitrate** | Reduced based on server boost level. |
|
||||
| **Image Limit** | Max 9 images per message/gallery. Images >10MB may be refused. |
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
1. **Videos not playing in Discord**: Try h264-MP4 format, reduce quality, enable `add_time`.
|
||||
2. **"ffmpeg not found"**: Install system ffmpeg and restart ComfyUI.
|
||||
3. **Webhook errors**: Verify URL validity and channel existence.
|
||||
4. **UUID Conflicts**: Ensure unique output directories for multiple users.
|
||||
5. **GitHub Issues**: Check token permissions (`public_repo`) and repo access.
|
||||
|
||||
> [!WARNING]
|
||||
> **Critical Issue**: Do NOT disable `add_time` for videos sent to Discord. They may appear as a single frame due to Discord's processing. Disable `include_video_info` instead if you want to hide details.
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
Contributions are welcome! Please see the [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to get started. Whether it's bug reports, feature suggestions, or pull requests, your help is appreciated.
|
||||
|
||||
1. Fork the Project
|
||||
2. Create your Feature Branch (`git checkout -b feature/AmazingFeature`)
|
||||
3. Commit your Changes (`git commit -m 'Add some AmazingFeature'`)
|
||||
4. Push to the Branch (`git push origin feature/AmazingFeature`)
|
||||
5. Open a Pull Request
|
||||
|
||||
## 📜 License
|
||||
|
||||
This project is licensed under the [GPL-3.0 License](LICENSE.md) - see the LICENSE file for details.
|
||||
|
||||
---
|
||||
|
||||
<div align="center">
|
||||
<br/>
|
||||
<h3>Developed by <a href="https://aemotionstudio.org/">Æmotion Studio</a></h3>
|
||||
|
||||
<a href="https://www.youtube.com/@aemotionstudio/videos">
|
||||
<img src="https://img.shields.io/badge/YouTube-FF0000?style=for-the-badge&logo=youtube&logoColor=white" alt="YouTube" />
|
||||
</a>
|
||||
<a href="https://discord.gg/UzC9353mfp">
|
||||
<img src="https://img.shields.io/badge/Discord-5865F2?style=for-the-badge&logo=discord&logoColor=white" alt="Discord" />
|
||||
</a>
|
||||
<a href="https://ko-fi.com/aemotionstudio">
|
||||
<img src="https://img.shields.io/badge/Ko--fi-F16061?style=for-the-badge&logo=ko-fi&logoColor=white" alt="Ko-fi" />
|
||||
</a>
|
||||
|
||||
<br/>
|
||||
<br/>
|
||||
|
||||
<p>Happy Creating! 🎨</p>
|
||||
</div>
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
"""ComfyUI nodes for sending images and video to Discord and saving them locally"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Get the current directory
|
||||
current_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
# Add to path if not already there
|
||||
if current_dir not in sys.path:
|
||||
sys.path.insert(0, current_dir)
|
||||
|
||||
# Import nodes from nodes package using relative imports
|
||||
from .nodes.image_node import DiscordSendSaveImage
|
||||
from .nodes.video_node import DiscordSendSaveVideo
|
||||
|
||||
# Node class mappings for ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DiscordSendSaveImage": DiscordSendSaveImage,
|
||||
"DiscordSendSaveVideo": DiscordSendSaveVideo
|
||||
}
|
||||
|
||||
# Display names for ComfyUI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DiscordSendSaveImage": "Discord Send & Save Image 📤",
|
||||
"DiscordSendSaveVideo": "Discord Send & Save Video 🎬"
|
||||
}
|
||||
|
||||
# Export all required elements
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
ComfyUI-DiscordSend Companion Bot
|
||||
|
||||
A Discord bot for triggering ComfyUI generations, managing queues,
|
||||
saving prompt templates, and browsing generation history.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add parent directory to path for importing shared utils
|
||||
_parent_dir = Path(__file__).parent.parent
|
||||
if str(_parent_dir) not in sys.path:
|
||||
sys.path.insert(0, str(_parent_dir))
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__author__ = "AEmotionStudio"
|
||||
@@ -0,0 +1,42 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Add project root to python path to allow imports
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
sys.path.append(str(project_root))
|
||||
|
||||
from bot.config import BotConfig
|
||||
from bot.bot import ComfyUIBot
|
||||
from shared.logging_config import setup_logging
|
||||
|
||||
def main():
|
||||
# Setup logging
|
||||
setup_logging()
|
||||
logger = logging.getLogger("bot")
|
||||
|
||||
# Load configuration
|
||||
try:
|
||||
config = BotConfig.load()
|
||||
except Exception as e:
|
||||
logger.critical(f"Failed to load configuration: {e}")
|
||||
return
|
||||
|
||||
if not config.discord.token:
|
||||
logger.critical("Discord token not found! Set DISCORDBOT_DISCORD_TOKEN env var or config.yaml")
|
||||
return
|
||||
|
||||
# Initialize and run bot
|
||||
bot = ComfyUIBot(config)
|
||||
|
||||
try:
|
||||
bot.run(config.discord.token)
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Bot stopped by user.")
|
||||
except Exception as e:
|
||||
logger.critical(f"Bot crashed: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+126
@@ -0,0 +1,126 @@
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
import sys
|
||||
import asyncio
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from .config import BotConfig
|
||||
from .database.repository import Repository
|
||||
from .comfyui.client import ComfyUIClient
|
||||
from .comfyui.websocket import ComfyUIWebSocket
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ComfyUIBot(commands.Bot):
|
||||
"""
|
||||
Main Bot Class for ComfyUI Companion.
|
||||
"""
|
||||
|
||||
def __init__(self, config: BotConfig):
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True # Needed for some commands if not pure slash
|
||||
intents.members = True # Useful for permission checks
|
||||
|
||||
super().__init__(
|
||||
command_prefix=commands.when_mentioned,
|
||||
intents=intents,
|
||||
help_command=None,
|
||||
description="ComfyUI Companion Bot"
|
||||
)
|
||||
self.config = config
|
||||
|
||||
# Database
|
||||
self.repository = Repository(config.database.url)
|
||||
|
||||
self.client_id = str(uuid.uuid4())
|
||||
|
||||
# ComfyUI Clients
|
||||
self.comfy_client = ComfyUIClient(base_url=config.comfyui.url)
|
||||
self.comfy_ws = ComfyUIWebSocket(base_url=config.comfyui.url, client_id=self.client_id)
|
||||
|
||||
# Services
|
||||
from .services.delivery import DeliveryService
|
||||
from .services.job_manager import JobManager
|
||||
from .services.permissions import PermissionService
|
||||
|
||||
self.delivery_service = DeliveryService(self, self.comfy_client)
|
||||
self.job_manager = JobManager(
|
||||
self.repository,
|
||||
self.comfy_client,
|
||||
self.comfy_ws,
|
||||
self.delivery_service
|
||||
)
|
||||
self.permission_service = PermissionService(self.repository)
|
||||
|
||||
async def setup_hook(self):
|
||||
"""Async setup before bot starts."""
|
||||
logger.info(f"Setting up bot for {self.user}...")
|
||||
|
||||
# Initialize Database
|
||||
await self.repository.init_db()
|
||||
logger.info("Database initialized.")
|
||||
|
||||
# Connect to ComfyUI WebSocket
|
||||
try:
|
||||
await self.comfy_ws.connect()
|
||||
await self.job_manager.start()
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not connect to ComfyUI WebSocket at start: {e}")
|
||||
# We don't crash, as ComfyUI might come up later
|
||||
|
||||
# Load Cogs
|
||||
await self._load_cogs()
|
||||
|
||||
# Sync Commands
|
||||
# In production, sync might be manual or per-guild to avoid rate limits
|
||||
# For simplicity in this self-hosted bot, we'll sync global
|
||||
try:
|
||||
synced = await self.tree.sync()
|
||||
logger.info(f"Synced {len(synced)} slash commands.")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to sync commands: {e}")
|
||||
|
||||
async def _load_cogs(self):
|
||||
"""Load extensions from cogs directory."""
|
||||
# We need to use dotted path: bot.cogs.generate
|
||||
cogs_dir = Path(__file__).parent / "cogs"
|
||||
|
||||
# Extensions to load
|
||||
extensions = [
|
||||
"bot.cogs.generate",
|
||||
"bot.cogs.queue",
|
||||
"bot.cogs.templates",
|
||||
"bot.cogs.history",
|
||||
"bot.cogs.admin",
|
||||
]
|
||||
|
||||
for ext in extensions:
|
||||
try:
|
||||
# Check if file exists first to avoid confusing errors if we haven't created it yet
|
||||
# (Since we are building incrementally)
|
||||
module_name = ext.split(".")[-1]
|
||||
if (cogs_dir / f"{module_name}.py").exists():
|
||||
await self.load_extension(ext)
|
||||
logger.info(f"Loaded extension: {ext}")
|
||||
else:
|
||||
logger.debug(f"Skipping extension {ext} (file not found)")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load extension {ext}: {e}")
|
||||
|
||||
async def close(self):
|
||||
"""Cleanup on shutdown."""
|
||||
logger.info("Shutting down bot...")
|
||||
await self.comfy_client.close()
|
||||
await self.comfy_ws.disconnect()
|
||||
await self.repository.close()
|
||||
await super().close()
|
||||
|
||||
async def on_ready(self):
|
||||
logger.info(f"Bot logged in as {self.user} (ID: {self.user.id})")
|
||||
logger.info(f"Connected to {len(self.guilds)} guilds")
|
||||
|
||||
async def on_command_error(self, ctx, error):
|
||||
"""Global error handler for prefix commands (if any)."""
|
||||
logger.error(f"Command error: {error}", exc_info=False)
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Discord bot cogs (command groups)."""
|
||||
|
||||
from .generate import GenerateCog
|
||||
from .queue import QueueCog
|
||||
from .templates import TemplateCog
|
||||
from .history import HistoryCog
|
||||
from .admin import AdminCog
|
||||
|
||||
__all__ = [
|
||||
"GenerateCog",
|
||||
"QueueCog",
|
||||
"TemplateCog",
|
||||
"HistoryCog",
|
||||
"AdminCog",
|
||||
]
|
||||
@@ -0,0 +1,68 @@
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
|
||||
from ..services.permissions import require_permission, Permissions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class AdminCog(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
@app_commands.command(name="admin", description="Bot configuration (Admin only)")
|
||||
@require_permission(Permissions.ADMIN.value)
|
||||
@app_commands.choices(action=[
|
||||
app_commands.Choice(name="status", value="status"),
|
||||
])
|
||||
async def admin(self, interaction: discord.Interaction, action: app_commands.Choice[str]):
|
||||
if action.value == "status":
|
||||
await self._show_status(interaction)
|
||||
|
||||
async def _show_status(self, interaction: discord.Interaction):
|
||||
comfy_status = await self.bot.comfy_client.check_status()
|
||||
status_emoji = "✅" if comfy_status else "❌"
|
||||
|
||||
embed = discord.Embed(title="Bot Status", color=discord.Color.dark_grey())
|
||||
embed.add_field(name="ComfyUI Connection", value=f"{status_emoji} {self.bot.config.comfyui.url}", inline=False)
|
||||
embed.add_field(name="Guilds", value=str(len(self.bot.guilds)), inline=True)
|
||||
embed.add_field(name="Latency", value=f"{round(self.bot.latency * 1000)}ms", inline=True)
|
||||
|
||||
await interaction.response.send_message(embed=embed, ephemeral=True)
|
||||
|
||||
@app_commands.command(name="setrole", description="Set permission level for a role")
|
||||
@require_permission(Permissions.ADMIN.value)
|
||||
@app_commands.choices(level=[
|
||||
app_commands.Choice(name="User", value="user"),
|
||||
app_commands.Choice(name="Generator", value="generator"),
|
||||
app_commands.Choice(name="Admin", value="admin"),
|
||||
])
|
||||
async def setrole(self, interaction: discord.Interaction, role: discord.Role, level: app_commands.Choice[str]):
|
||||
"""Assign a permission level to a Discord role."""
|
||||
if not interaction.guild:
|
||||
await interaction.response.send_message("This command must be used in a server.", ephemeral=True)
|
||||
return
|
||||
|
||||
try:
|
||||
# Ensure server exists in DB
|
||||
await self.bot.repository.get_or_create_server(
|
||||
str(interaction.guild.id),
|
||||
interaction.guild.name
|
||||
)
|
||||
|
||||
await self.bot.repository.set_server_role(
|
||||
server_discord_id=str(interaction.guild.id),
|
||||
role_discord_id=str(role.id),
|
||||
permission_level=level.value
|
||||
)
|
||||
await interaction.response.send_message(
|
||||
f"✅ Role {role.mention} set to **{level.name}** permission level.",
|
||||
ephemeral=True
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to set role: {e}")
|
||||
await interaction.response.send_message("❌ Failed to update role permissions.", ephemeral=True)
|
||||
|
||||
async def setup(bot):
|
||||
await bot.add_cog(AdminCog(bot))
|
||||
@@ -0,0 +1,120 @@
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ..embeds.builders import EmbedBuilder
|
||||
from ..services.permissions import require_permission, Permissions
|
||||
from shared.workflow import WorkflowBuilder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class GenerateCog(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
@app_commands.command(name="generate", description="Generate an image using ComfyUI")
|
||||
@app_commands.describe(
|
||||
prompt="The positive prompt to generate",
|
||||
negative_prompt="Aspects to avoid (optional)",
|
||||
seed="Seed for generation (optional)",
|
||||
steps="Number of steps (optional)",
|
||||
cfg="CFG Scale (optional)",
|
||||
delivery="Delivery method (channel or dm)"
|
||||
)
|
||||
@app_commands.choices(delivery=[
|
||||
app_commands.Choice(name="Current Channel", value="channel"),
|
||||
app_commands.Choice(name="Direct Message", value="dm")
|
||||
])
|
||||
@require_permission(Permissions.USER.value)
|
||||
async def generate(self, interaction: discord.Interaction,
|
||||
prompt: str,
|
||||
negative_prompt: str = "",
|
||||
seed: int = None,
|
||||
steps: int = None,
|
||||
cfg: float = None,
|
||||
delivery: app_commands.Choice[str] = None):
|
||||
|
||||
await interaction.response.defer()
|
||||
|
||||
# Load default workflow
|
||||
# TODO: Move this path to config or database
|
||||
workflow_path = Path(__file__).parent.parent / "data" / "default_workflow_api.json"
|
||||
|
||||
if not workflow_path.exists():
|
||||
await interaction.followup.send("❌ Error: Default workflow not found.", ephemeral=True)
|
||||
return
|
||||
|
||||
try:
|
||||
with open(workflow_path, "r") as f:
|
||||
workflow_json = json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load workflow: {e}")
|
||||
await interaction.followup.send("❌ Error: Failed to load workflow configuration.", ephemeral=True)
|
||||
return
|
||||
|
||||
# Prepare parameters for tracking
|
||||
parameters = {
|
||||
"seed": seed,
|
||||
"steps": steps,
|
||||
"cfg": cfg
|
||||
}
|
||||
|
||||
# Modify workflow
|
||||
builder = WorkflowBuilder(workflow_json)
|
||||
builder.set_prompt(prompt, negative_prompt)
|
||||
|
||||
if seed is not None:
|
||||
builder.set_seed(seed)
|
||||
else:
|
||||
# Random seed if not provided
|
||||
import random
|
||||
generated_seed = random.randint(1, 1000000000000000)
|
||||
builder.set_seed(generated_seed)
|
||||
parameters["seed"] = generated_seed # Track actual seed
|
||||
|
||||
if steps is not None:
|
||||
builder.set_steps(steps)
|
||||
|
||||
if cfg is not None:
|
||||
builder.set_cfg(cfg)
|
||||
|
||||
final_workflow = builder.get_workflow()
|
||||
|
||||
# Determine delivery method
|
||||
delivery_method = delivery.value if delivery else "channel"
|
||||
|
||||
# Check server context
|
||||
server_id = str(interaction.guild_id) if interaction.guild else None
|
||||
channel_id = str(interaction.channel_id)
|
||||
|
||||
try:
|
||||
# Create Job
|
||||
job = await self.bot.job_manager.create_job(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
workflow=final_workflow,
|
||||
positive_prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
parameters=parameters,
|
||||
server_discord_id=server_id,
|
||||
channel_id=channel_id,
|
||||
delivery_type=delivery_method
|
||||
)
|
||||
|
||||
# Send Queued Embed
|
||||
embed = EmbedBuilder.job_queued(job)
|
||||
await interaction.followup.send(embed=embed)
|
||||
|
||||
# Store the interaction message ID if we want to update it later
|
||||
# (JobManager could use this to update the specific message)
|
||||
original_message = await interaction.original_response()
|
||||
await self.bot.repository.update_job_message(job.prompt_id, str(original_message.id))
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start generation: {e}")
|
||||
await interaction.followup.send(f"❌ Error starting generation: {str(e)}", ephemeral=True)
|
||||
|
||||
async def setup(bot):
|
||||
await bot.add_cog(GenerateCog(bot))
|
||||
@@ -0,0 +1,209 @@
|
||||
"""History cog for viewing past generations and rerunning them."""
|
||||
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
import json
|
||||
import random
|
||||
from typing import List
|
||||
|
||||
from ..services.permissions import require_permission, Permissions
|
||||
from ..database.models import JobStatus
|
||||
from ..embeds.builders import EmbedBuilder
|
||||
from shared.workflow import WorkflowBuilder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HistoryPaginator(discord.ui.View):
|
||||
"""Paginated view for job history."""
|
||||
|
||||
def __init__(self, jobs: List, per_page: int = 5):
|
||||
super().__init__(timeout=180)
|
||||
self.jobs = jobs
|
||||
self.per_page = per_page
|
||||
self.page = 0
|
||||
self.max_page = (len(jobs) - 1) // per_page if jobs else 0
|
||||
self._update_buttons()
|
||||
|
||||
def _update_buttons(self):
|
||||
self.prev_button.disabled = self.page <= 0
|
||||
self.next_button.disabled = self.page >= self.max_page
|
||||
|
||||
def get_embed(self) -> discord.Embed:
|
||||
embed = discord.Embed(title="Generation History", color=discord.Color.blue())
|
||||
|
||||
start = self.page * self.per_page
|
||||
end = start + self.per_page
|
||||
page_jobs = self.jobs[start:end]
|
||||
|
||||
if not page_jobs:
|
||||
embed.description = "No generation history found."
|
||||
return embed
|
||||
|
||||
lines = []
|
||||
for job in page_jobs:
|
||||
status_emoji = {
|
||||
JobStatus.COMPLETED.value: "✅",
|
||||
JobStatus.FAILED.value: "❌",
|
||||
JobStatus.CANCELLED.value: "🚫",
|
||||
JobStatus.PENDING.value: "⏳",
|
||||
JobStatus.RUNNING.value: "🔄",
|
||||
}.get(job.status, "❓")
|
||||
|
||||
prompt_preview = (job.positive_prompt or "No prompt")[:50]
|
||||
if len(job.positive_prompt or "") > 50:
|
||||
prompt_preview += "..."
|
||||
|
||||
timestamp = (
|
||||
job.created_at.strftime("%Y-%m-%d %H:%M") if job.created_at else "Unknown"
|
||||
)
|
||||
|
||||
lines.append(
|
||||
f"{status_emoji} **ID: {job.id}** | {timestamp}\n└ {prompt_preview}"
|
||||
)
|
||||
|
||||
embed.description = "\n\n".join(lines)
|
||||
embed.set_footer(
|
||||
text=f"Page {self.page + 1}/{self.max_page + 1} | Use /rerun <id> to regenerate"
|
||||
)
|
||||
|
||||
return embed
|
||||
|
||||
@discord.ui.button(label="◀ Previous", style=discord.ButtonStyle.secondary)
|
||||
async def prev_button(
|
||||
self, interaction: discord.Interaction, button: discord.ui.Button
|
||||
):
|
||||
self.page = max(0, self.page - 1)
|
||||
self._update_buttons()
|
||||
await interaction.response.edit_message(embed=self.get_embed(), view=self)
|
||||
|
||||
@discord.ui.button(label="Next ▶", style=discord.ButtonStyle.secondary)
|
||||
async def next_button(
|
||||
self, interaction: discord.Interaction, button: discord.ui.Button
|
||||
):
|
||||
self.page = min(self.max_page, self.page + 1)
|
||||
self._update_buttons()
|
||||
await interaction.response.edit_message(embed=self.get_embed(), view=self)
|
||||
|
||||
|
||||
class HistoryCog(commands.Cog):
|
||||
"""View generation history and rerun past jobs."""
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
@app_commands.command(name="history", description="View your generation history")
|
||||
@app_commands.describe(limit="Number of jobs to show (default: 20, max: 50)")
|
||||
async def history(self, interaction: discord.Interaction, limit: int = 20):
|
||||
"""Show paginated generation history."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
limit = min(max(1, limit), 50) # Clamp between 1 and 50
|
||||
|
||||
jobs = await self.bot.repository.list_user_jobs(
|
||||
user_discord_id=str(interaction.user.id), limit=limit
|
||||
)
|
||||
|
||||
if not jobs:
|
||||
await interaction.followup.send(
|
||||
"You have no generation history.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
view = HistoryPaginator(jobs)
|
||||
await interaction.followup.send(embed=view.get_embed(), view=view, ephemeral=True)
|
||||
|
||||
@app_commands.command(name="rerun", description="Rerun a previous generation")
|
||||
@app_commands.describe(job_id="The ID of the job to rerun")
|
||||
@require_permission(Permissions.USER.value)
|
||||
async def rerun(self, interaction: discord.Interaction, job_id: int):
|
||||
"""Rerun a previous job with the same parameters."""
|
||||
await interaction.response.defer()
|
||||
|
||||
# Get the original job
|
||||
original_job = await self.bot.repository.get_job_by_id(job_id)
|
||||
|
||||
if not original_job:
|
||||
await interaction.followup.send(
|
||||
f"Job ID {job_id} not found.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
# Check ownership
|
||||
if str(original_job.user.discord_id) != str(interaction.user.id):
|
||||
await interaction.followup.send(
|
||||
"You can only rerun your own jobs.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
# Check if job has required data
|
||||
if not original_job.workflow_json:
|
||||
await interaction.followup.send(
|
||||
"This job cannot be rerun (workflow data not saved).", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
workflow = json.loads(original_job.workflow_json)
|
||||
except json.JSONDecodeError:
|
||||
await interaction.followup.send(
|
||||
"Failed to parse original workflow.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
# Parse parameters
|
||||
parameters = {}
|
||||
if original_job.parameters:
|
||||
try:
|
||||
parameters = json.loads(original_job.parameters)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Generate new seed for rerun
|
||||
new_seed = random.randint(1, 1000000000000000)
|
||||
parameters["seed"] = new_seed
|
||||
|
||||
# Update workflow with new seed
|
||||
builder = WorkflowBuilder(workflow)
|
||||
builder.set_seed(new_seed)
|
||||
final_workflow = builder.get_workflow()
|
||||
|
||||
# Determine delivery and context
|
||||
server_id = str(interaction.guild_id) if interaction.guild else None
|
||||
channel_id = str(interaction.channel_id)
|
||||
|
||||
try:
|
||||
# Create new job
|
||||
job = await self.bot.job_manager.create_job(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
workflow=final_workflow,
|
||||
positive_prompt=original_job.positive_prompt or "",
|
||||
negative_prompt=original_job.negative_prompt or "",
|
||||
parameters=parameters,
|
||||
server_discord_id=server_id,
|
||||
channel_id=channel_id,
|
||||
delivery_type=original_job.delivery_type or "channel",
|
||||
)
|
||||
|
||||
embed = EmbedBuilder.job_queued(job)
|
||||
embed.set_footer(text=f"Rerun of job #{job_id} | New job ID: {job.id}")
|
||||
|
||||
await interaction.followup.send(embed=embed)
|
||||
|
||||
# Store message ID
|
||||
original_message = await interaction.original_response()
|
||||
await self.bot.repository.update_job_message(
|
||||
job.prompt_id, str(original_message.id)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to rerun job: {e}")
|
||||
await interaction.followup.send(
|
||||
f"Failed to rerun job: {str(e)}", ephemeral=True
|
||||
)
|
||||
|
||||
|
||||
async def setup(bot):
|
||||
await bot.add_cog(HistoryCog(bot))
|
||||
@@ -0,0 +1,115 @@
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from ..services.permissions import require_permission, Permissions
|
||||
from ..database.models import JobStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class QueueCog(commands.Cog):
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
@app_commands.command(name="queue", description="Manage the generation queue")
|
||||
@app_commands.choices(action=[
|
||||
app_commands.Choice(name="view", value="view"),
|
||||
app_commands.Choice(name="clear", value="clear"),
|
||||
])
|
||||
async def queue(self, interaction: discord.Interaction, action: app_commands.Choice[str]):
|
||||
"""General queue commands."""
|
||||
command = action.value
|
||||
|
||||
if command == "view":
|
||||
await self._view_queue(interaction)
|
||||
elif command == "clear":
|
||||
await self._clear_queue(interaction)
|
||||
|
||||
async def _view_queue(self, interaction: discord.Interaction):
|
||||
"""Show current pending jobs."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
pending_jobs = await self.bot.repository.get_pending_jobs()
|
||||
|
||||
if not pending_jobs:
|
||||
await interaction.followup.send("🟢 The queue is currently empty.", ephemeral=True)
|
||||
return
|
||||
|
||||
embed = discord.Embed(
|
||||
title=f"Generation Queue ({len(pending_jobs)})",
|
||||
color=discord.Color.blue()
|
||||
)
|
||||
|
||||
# Show top 10
|
||||
desc_lines = []
|
||||
for i, job in enumerate(pending_jobs[:10]):
|
||||
status_icon = "🔄" if job.status == JobStatus.RUNNING.value else "⏳"
|
||||
user_mention = f"<@{job.user.discord_id}>"
|
||||
prompt_text = job.positive_prompt or "No prompt provided"
|
||||
prompt_preview = (prompt_text[:40] + "...") if len(prompt_text) > 40 else prompt_text
|
||||
desc_lines.append(f"`#{i+1}` {status_icon} **ID:{job.id}** {user_mention}: {prompt_preview}")
|
||||
|
||||
if len(pending_jobs) > 10:
|
||||
desc_lines.append(f"...and {len(pending_jobs) - 10} more.")
|
||||
|
||||
embed.description = "\n".join(desc_lines)
|
||||
await interaction.followup.send(embed=embed, ephemeral=True)
|
||||
|
||||
async def _clear_queue(self, interaction: discord.Interaction):
|
||||
"""Clear user's own pending jobs."""
|
||||
# Check permission manually since it's a subcommand handler
|
||||
has_perm = await self.bot.permission_service.check_permission(interaction.user, Permissions.GENERATOR.value)
|
||||
if not has_perm:
|
||||
await interaction.response.send_message("⛔ You need the **Generator** role to clear the queue.", ephemeral=True)
|
||||
return
|
||||
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
# Get user's pending jobs
|
||||
# Note: repo.get_pending_jobs returns ALL. Better to filter or add new repo method.
|
||||
# But for now let's iterate.
|
||||
all_pending = await self.bot.repository.get_pending_jobs()
|
||||
user_jobs = [j for j in all_pending if str(j.user.discord_id) == str(interaction.user.id)]
|
||||
|
||||
count = 0
|
||||
for job in user_jobs:
|
||||
success = await self.bot.job_manager.cancel_job(job.id)
|
||||
if success:
|
||||
count += 1
|
||||
|
||||
if count > 0:
|
||||
await interaction.followup.send(f"🗑️ Cancelled {count} of your pending jobs.", ephemeral=True)
|
||||
else:
|
||||
await interaction.followup.send("No pending jobs found to clear.", ephemeral=True)
|
||||
|
||||
|
||||
@app_commands.command(name="cancel", description="Cancel a specific job")
|
||||
@app_commands.describe(job_id="The ID of the job to cancel")
|
||||
@require_permission(Permissions.GENERATOR.value)
|
||||
async def cancel(self, interaction: discord.Interaction, job_id: int):
|
||||
"""Cancel a specific job by ID."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
job = await self.bot.repository.get_job_by_id(job_id)
|
||||
if not job:
|
||||
await interaction.followup.send(f"❌ Job ID {job_id} not found.", ephemeral=True)
|
||||
return
|
||||
|
||||
# Check permissions
|
||||
is_owner = str(job.user.discord_id) == str(interaction.user.id)
|
||||
is_admin = await self.bot.permission_service.check_permission(interaction.user, Permissions.ADMIN.value)
|
||||
|
||||
if not is_owner and not is_admin:
|
||||
await interaction.followup.send("⛔ You can only cancel your own jobs.", ephemeral=True)
|
||||
return
|
||||
|
||||
success = await self.bot.job_manager.cancel_job(job_id)
|
||||
if success:
|
||||
await interaction.followup.send(f"✅ Job {job_id} cancelled.", ephemeral=True)
|
||||
else:
|
||||
await interaction.followup.send(f"⚠️ Could not cancel job {job_id} (maybe already finished?).", ephemeral=True)
|
||||
|
||||
async def setup(bot):
|
||||
await bot.add_cog(QueueCog(bot))
|
||||
@@ -0,0 +1,223 @@
|
||||
"""Template management cog for saving and loading prompt presets."""
|
||||
|
||||
import discord
|
||||
from discord import app_commands
|
||||
from discord.ext import commands
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
from ..services.permissions import require_permission, Permissions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TemplateCog(commands.Cog):
|
||||
"""Manage prompt templates."""
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
template_group = app_commands.Group(
|
||||
name="template", description="Manage prompt templates"
|
||||
)
|
||||
|
||||
@template_group.command(name="save", description="Save a prompt as a template")
|
||||
@app_commands.describe(
|
||||
name="Template name (unique per user/server)",
|
||||
prompt="The positive prompt to save",
|
||||
negative_prompt="Negative prompt (optional)",
|
||||
shared="Share with entire server (default: private)",
|
||||
)
|
||||
@require_permission(Permissions.GENERATOR.value)
|
||||
async def template_save(
|
||||
self,
|
||||
interaction: discord.Interaction,
|
||||
name: str,
|
||||
prompt: str,
|
||||
negative_prompt: str = "",
|
||||
shared: bool = False,
|
||||
):
|
||||
"""Save current prompt as a named template."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
# Validate name
|
||||
if len(name) > 100:
|
||||
await interaction.followup.send(
|
||||
"Template name must be 100 characters or less.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
# Ensure user exists
|
||||
await self.bot.repository.get_or_create_user(
|
||||
str(interaction.user.id), interaction.user.display_name
|
||||
)
|
||||
|
||||
server_id = str(interaction.guild_id) if interaction.guild and shared else None
|
||||
if server_id:
|
||||
await self.bot.repository.get_or_create_server(
|
||||
server_id, interaction.guild.name
|
||||
)
|
||||
|
||||
try:
|
||||
await self.bot.repository.create_template(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
name=name,
|
||||
positive_prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
server_discord_id=server_id,
|
||||
)
|
||||
|
||||
scope = "server" if shared else "private"
|
||||
await interaction.followup.send(
|
||||
f"Saved template **{name}** ({scope}).", ephemeral=True
|
||||
)
|
||||
except Exception as e:
|
||||
if "UNIQUE constraint" in str(e):
|
||||
await interaction.followup.send(
|
||||
f"A template named **{name}** already exists. "
|
||||
"Delete it first or use a different name.",
|
||||
ephemeral=True,
|
||||
)
|
||||
else:
|
||||
logger.error(f"Failed to save template: {e}")
|
||||
await interaction.followup.send(
|
||||
"Failed to save template.", ephemeral=True
|
||||
)
|
||||
|
||||
@template_group.command(name="load", description="Load a saved template")
|
||||
@app_commands.describe(name="Template name to load")
|
||||
async def template_load(self, interaction: discord.Interaction, name: str):
|
||||
"""Load a template and show its contents."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
server_id = str(interaction.guild_id) if interaction.guild else None
|
||||
|
||||
# Try user's private template first
|
||||
template = await self.bot.repository.get_template(
|
||||
user_discord_id=str(interaction.user.id), name=name
|
||||
)
|
||||
|
||||
# Try shared server template if not found
|
||||
if not template and server_id:
|
||||
templates = await self.bot.repository.list_templates(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
server_discord_id=server_id,
|
||||
include_shared=True,
|
||||
)
|
||||
template = next((t for t in templates if t.name == name), None)
|
||||
|
||||
if not template:
|
||||
await interaction.followup.send(
|
||||
f"Template **{name}** not found.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
embed = discord.Embed(title=f"Template: {template.name}", color=discord.Color.blue())
|
||||
embed.add_field(
|
||||
name="Prompt", value=template.positive_prompt[:1024], inline=False
|
||||
)
|
||||
if template.negative_prompt:
|
||||
embed.add_field(
|
||||
name="Negative Prompt",
|
||||
value=template.negative_prompt[:1024],
|
||||
inline=False,
|
||||
)
|
||||
|
||||
scope = "Shared" if template.server_id else "Private"
|
||||
embed.set_footer(text=f"{scope} template | Use /generate with this prompt")
|
||||
|
||||
await interaction.followup.send(embed=embed, ephemeral=True)
|
||||
|
||||
@template_group.command(name="list", description="List your saved templates")
|
||||
async def template_list(self, interaction: discord.Interaction):
|
||||
"""List all available templates."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
server_id = str(interaction.guild_id) if interaction.guild else None
|
||||
|
||||
templates = await self.bot.repository.list_templates(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
server_discord_id=server_id,
|
||||
include_shared=True,
|
||||
)
|
||||
|
||||
if not templates:
|
||||
await interaction.followup.send(
|
||||
"You have no saved templates.", ephemeral=True
|
||||
)
|
||||
return
|
||||
|
||||
embed = discord.Embed(title="Your Templates", color=discord.Color.blue())
|
||||
|
||||
private_templates = [t for t in templates if t.server_id is None]
|
||||
shared_templates = [t for t in templates if t.server_id is not None]
|
||||
|
||||
if private_templates:
|
||||
names = "\n".join([f"• {t.name}" for t in private_templates[:10]])
|
||||
if len(private_templates) > 10:
|
||||
names += f"\n...and {len(private_templates) - 10} more"
|
||||
embed.add_field(name="Private Templates", value=names, inline=False)
|
||||
|
||||
if shared_templates:
|
||||
names = "\n".join([f"• {t.name}" for t in shared_templates[:10]])
|
||||
if len(shared_templates) > 10:
|
||||
names += f"\n...and {len(shared_templates) - 10} more"
|
||||
embed.add_field(name="Server Templates", value=names, inline=False)
|
||||
|
||||
await interaction.followup.send(embed=embed, ephemeral=True)
|
||||
|
||||
@template_group.command(name="delete", description="Delete a saved template")
|
||||
@app_commands.describe(name="Template name to delete")
|
||||
@require_permission(Permissions.GENERATOR.value)
|
||||
async def template_delete(self, interaction: discord.Interaction, name: str):
|
||||
"""Delete a template."""
|
||||
await interaction.response.defer(ephemeral=True)
|
||||
|
||||
# Try deleting private template
|
||||
deleted = await self.bot.repository.delete_template(
|
||||
user_discord_id=str(interaction.user.id), name=name
|
||||
)
|
||||
|
||||
# Try deleting shared template if private not found
|
||||
if not deleted and interaction.guild:
|
||||
deleted = await self.bot.repository.delete_template(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
name=name,
|
||||
server_discord_id=str(interaction.guild_id),
|
||||
)
|
||||
|
||||
if deleted:
|
||||
await interaction.followup.send(
|
||||
f"Deleted template **{name}**.", ephemeral=True
|
||||
)
|
||||
else:
|
||||
await interaction.followup.send(
|
||||
f"Template **{name}** not found or you don't have permission to delete it.",
|
||||
ephemeral=True,
|
||||
)
|
||||
|
||||
@template_load.autocomplete("name")
|
||||
@template_delete.autocomplete("name")
|
||||
async def template_name_autocomplete(
|
||||
self, interaction: discord.Interaction, current: str
|
||||
) -> List[app_commands.Choice[str]]:
|
||||
"""Autocomplete for template names."""
|
||||
server_id = str(interaction.guild_id) if interaction.guild else None
|
||||
|
||||
templates = await self.bot.repository.list_templates(
|
||||
user_discord_id=str(interaction.user.id),
|
||||
server_discord_id=server_id,
|
||||
include_shared=True,
|
||||
)
|
||||
|
||||
# Filter by current input
|
||||
filtered = [t for t in templates if current.lower() in t.name.lower()]
|
||||
|
||||
return [
|
||||
app_commands.Choice(name=t.name, value=t.name)
|
||||
for t in filtered[:25] # Discord limit
|
||||
]
|
||||
|
||||
|
||||
async def setup(bot):
|
||||
await bot.add_cog(TemplateCog(bot))
|
||||
@@ -0,0 +1,6 @@
|
||||
"""ComfyUI integration package."""
|
||||
|
||||
from .client import ComfyUIClient
|
||||
from .websocket import ComfyUIWebSocket
|
||||
|
||||
__all__ = ["ComfyUIClient", "ComfyUIWebSocket"]
|
||||
@@ -0,0 +1,107 @@
|
||||
import aiohttp
|
||||
import logging
|
||||
from typing import Optional, Dict, List, Any
|
||||
import json
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class ComfyUIClient:
|
||||
"""Async Client for interacting with ComfyUI REST API."""
|
||||
|
||||
def __init__(self, base_url: str = "http://127.0.0.1:8188"):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.session: Optional[aiohttp.ClientSession] = None
|
||||
|
||||
async def _get_session(self) -> aiohttp.ClientSession:
|
||||
if self.session is None or self.session.closed:
|
||||
self.session = aiohttp.ClientSession()
|
||||
return self.session
|
||||
|
||||
async def close(self):
|
||||
if self.session and not self.session.closed:
|
||||
await self.session.close()
|
||||
|
||||
async def check_status(self) -> bool:
|
||||
"""Check if ComfyUI server is reachable."""
|
||||
try:
|
||||
session = await self._get_session()
|
||||
async with session.get(f"{self.base_url}/system_stats") as response:
|
||||
return response.status == 200
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to connect to ComfyUI: {e}")
|
||||
return False
|
||||
|
||||
async def get_system_stats(self) -> Dict[str, Any]:
|
||||
"""Get system statistics."""
|
||||
session = await self._get_session()
|
||||
async with session.get(f"{self.base_url}/system_stats") as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
async def get_queue(self) -> Dict[str, Any]:
|
||||
"""Get current queue status."""
|
||||
session = await self._get_session()
|
||||
async with session.get(f"{self.base_url}/queue") as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
async def get_history(self, prompt_id: str) -> Dict[str, Any]:
|
||||
"""Get history for a specific prompt ID."""
|
||||
session = await self._get_session()
|
||||
async with session.get(f"{self.base_url}/history/{prompt_id}") as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
async def queue_prompt(self, workflow: Dict[str, Any], client_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Queue a workflow for generation.
|
||||
|
||||
Args:
|
||||
workflow: The workflow JSON object (API format)
|
||||
client_id: Unique client ID for WebSocket correlation
|
||||
"""
|
||||
session = await self._get_session()
|
||||
payload = {
|
||||
"prompt": workflow,
|
||||
"client_id": client_id
|
||||
}
|
||||
async with session.post(f"{self.base_url}/prompt", json=payload) as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
async def interrupt(self):
|
||||
"""Interrupt currently executing prompt."""
|
||||
session = await self._get_session()
|
||||
async with session.post(f"{self.base_url}/interrupt") as response:
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to interrupt: {e}")
|
||||
|
||||
async def delete_queue_item(self, prompt_id: str):
|
||||
"""Remove an item from queue."""
|
||||
session = await self._get_session()
|
||||
payload = {"delete": [prompt_id]}
|
||||
async with session.post(f"{self.base_url}/queue", json=payload) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
async def get_image(self, filename: str, subfolder: str = "", type: str = "output") -> bytes:
|
||||
"""Fetch a generated image."""
|
||||
session = await self._get_session()
|
||||
params = {"filename": filename, "subfolder": subfolder, "type": type}
|
||||
async with session.get(f"{self.base_url}/view", params=params) as response:
|
||||
response.raise_for_status()
|
||||
return await response.read()
|
||||
|
||||
async def upload_image(self, image_data: bytes, filename: str, subfolder: str = ""):
|
||||
"""Upload an image to ComfyUI (input folder)."""
|
||||
session = await self._get_session()
|
||||
data = aiohttp.FormData()
|
||||
data.add_field("image", image_data, filename=filename)
|
||||
if subfolder:
|
||||
data.add_field("subfolder", subfolder)
|
||||
|
||||
async with session.post(f"{self.base_url}/upload/image", data=data) as response:
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
@@ -0,0 +1,205 @@
|
||||
import aiohttp
|
||||
import logging
|
||||
import json
|
||||
import asyncio
|
||||
import random
|
||||
from typing import Callable, Coroutine, Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Reconnection constants
|
||||
INITIAL_BACKOFF = 1.0 # Initial delay in seconds
|
||||
MAX_BACKOFF = 60.0 # Maximum delay cap
|
||||
BACKOFF_MULTIPLIER = 2 # Exponential multiplier
|
||||
JITTER_FACTOR = 0.1 # +/- 10% randomization
|
||||
|
||||
class ComfyUIWebSocket:
|
||||
"""WebSocket Client for real-time ComfyUI events."""
|
||||
|
||||
def __init__(self, base_url: str = "ws://127.0.0.1:8188", client_id: str = ""):
|
||||
# Convert http/https to ws/wss if needed
|
||||
if base_url.startswith("http://"):
|
||||
base_url = base_url.replace("http://", "ws://")
|
||||
elif base_url.startswith("https://"):
|
||||
base_url = base_url.replace("https://", "wss://")
|
||||
|
||||
self.ws_url = f"{base_url.rstrip('/')}/ws"
|
||||
self.client_id = client_id
|
||||
if client_id:
|
||||
self.ws_url += f"?clientId={client_id}"
|
||||
|
||||
self.ws: Optional[aiohttp.ClientWebSocketResponse] = None
|
||||
self.session: Optional[aiohttp.ClientSession] = None
|
||||
self._callbacks: Dict[str, List[Callable[[Dict[str, Any]], Coroutine[Any, Any, None]]]] = {}
|
||||
self._running = False
|
||||
self._listen_task: Optional[asyncio.Task] = None
|
||||
|
||||
# Reconnection state
|
||||
self._reconnect_attempts = 0
|
||||
self._should_reconnect = True
|
||||
self._reconnect_task: Optional[asyncio.Task] = None
|
||||
|
||||
# Lock to prevent race conditions between connect() and disconnect()
|
||||
self._state_lock = asyncio.Lock()
|
||||
|
||||
async def connect(self) -> bool:
|
||||
"""Connect to the WebSocket.
|
||||
|
||||
Returns:
|
||||
True if connection was established successfully, False if aborted.
|
||||
"""
|
||||
async with self._state_lock:
|
||||
# Check if disconnect was called - don't proceed if so
|
||||
if not self._should_reconnect:
|
||||
logger.info("Connect aborted: disconnect was requested")
|
||||
return False
|
||||
|
||||
if self.session is None or self.session.closed:
|
||||
self.session = aiohttp.ClientSession()
|
||||
|
||||
try:
|
||||
self.ws = await self.session.ws_connect(self.ws_url)
|
||||
# Re-check after await in case disconnect() was called during connection
|
||||
if not self._should_reconnect:
|
||||
logger.info("Connect aborted after ws_connect: disconnect was requested")
|
||||
if self.ws and not self.ws.closed:
|
||||
await self.ws.close()
|
||||
if self.session and not self.session.closed:
|
||||
await self.session.close()
|
||||
return False
|
||||
self._running = True
|
||||
self._reconnect_attempts = 0 # Reset on successful connection
|
||||
self._listen_task = asyncio.create_task(self._listen())
|
||||
logger.info(f"Connected to ComfyUI WebSocket at {self.ws_url}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to WebSocket: {e}")
|
||||
if self.session and not self.session.closed:
|
||||
await self.session.close()
|
||||
raise
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from WebSocket."""
|
||||
async with self._state_lock:
|
||||
self._should_reconnect = False # Prevent reconnection loop
|
||||
self._running = False
|
||||
|
||||
# Cancel reconnection task if running (outside lock to avoid deadlock)
|
||||
if self._reconnect_task and not self._reconnect_task.done():
|
||||
self._reconnect_task.cancel()
|
||||
try:
|
||||
await self._reconnect_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self.ws:
|
||||
await self.ws.close()
|
||||
if self.session:
|
||||
await self.session.close()
|
||||
if self._listen_task:
|
||||
try:
|
||||
await self._listen_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
def add_listener(self, event_type: str, callback: Callable[[Dict[str, Any]], Coroutine[Any, Any, None]]):
|
||||
"""Register a callback for an event type."""
|
||||
if event_type not in self._callbacks:
|
||||
self._callbacks[event_type] = []
|
||||
self._callbacks[event_type].append(callback)
|
||||
|
||||
async def _listen(self):
|
||||
"""Listen loop for incoming messages."""
|
||||
if not self.ws:
|
||||
return
|
||||
|
||||
try:
|
||||
async for msg in self.ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
data = json.loads(msg.data)
|
||||
event_type = data.get("type", "unknown")
|
||||
# Some messages pack content in 'data', others at top level
|
||||
# ComfyUI typically sends {type: "event_name", data: {...}, sid: "..."}
|
||||
|
||||
handlers = self._callbacks.get(event_type, [])
|
||||
for handler in handlers:
|
||||
try:
|
||||
await handler(data)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in WebSocket handler for {event_type}: {e}")
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Received invalid JSON: {msg.data}")
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
logger.error("WebSocket connection closed with error")
|
||||
break
|
||||
except Exception as e:
|
||||
if self._running:
|
||||
logger.error(f"WebSocket listener error: {e}")
|
||||
finally:
|
||||
if self._running and self._should_reconnect:
|
||||
logger.warning("WebSocket connection lost. Starting reconnection...")
|
||||
self._reconnect_task = asyncio.create_task(self._handle_disconnect())
|
||||
|
||||
def _calculate_backoff(self) -> float:
|
||||
"""Calculate backoff delay with exponential growth and jitter."""
|
||||
delay = INITIAL_BACKOFF * (BACKOFF_MULTIPLIER ** self._reconnect_attempts)
|
||||
delay = min(delay, MAX_BACKOFF)
|
||||
# Add jitter: +/- JITTER_FACTOR
|
||||
jitter = delay * JITTER_FACTOR * (2 * random.random() - 1)
|
||||
return delay + jitter
|
||||
|
||||
async def _handle_disconnect(self) -> None:
|
||||
"""Handle unexpected disconnection by attempting to reconnect."""
|
||||
self._running = False
|
||||
|
||||
# Close existing connections
|
||||
if self.ws and not self.ws.closed:
|
||||
try:
|
||||
await self.ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
if self.session and not self.session.closed:
|
||||
try:
|
||||
await self.session.close()
|
||||
except Exception:
|
||||
pass
|
||||
self.session = None
|
||||
self.ws = None
|
||||
|
||||
await self._reconnect_loop()
|
||||
|
||||
async def _reconnect_loop(self) -> None:
|
||||
"""Background task that handles reconnection with exponential backoff."""
|
||||
while self._should_reconnect:
|
||||
self._reconnect_attempts += 1
|
||||
backoff = self._calculate_backoff()
|
||||
|
||||
logger.info(
|
||||
f"Reconnection attempt {self._reconnect_attempts} "
|
||||
f"in {backoff:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
|
||||
if not self._should_reconnect:
|
||||
logger.info("Reconnection cancelled.")
|
||||
return
|
||||
|
||||
try:
|
||||
attempts_made = self._reconnect_attempts
|
||||
connected = await self.connect()
|
||||
if connected:
|
||||
logger.info(
|
||||
f"Successfully reconnected after "
|
||||
f"{attempts_made} attempt(s)."
|
||||
)
|
||||
return
|
||||
else:
|
||||
# connect() returned False - disconnect was called
|
||||
logger.info("Reconnection aborted: disconnect was requested")
|
||||
return
|
||||
except Exception as e:
|
||||
logger.warning(f"Reconnection attempt failed: {e}")
|
||||
|
||||
logger.info("Reconnection loop ended (should_reconnect=False).")
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
"""
|
||||
Configuration management for the Discord bot.
|
||||
|
||||
Configuration is loaded from (in priority order):
|
||||
1. Environment variables (highest priority)
|
||||
2. Config file (bot/config.yaml)
|
||||
3. Defaults (lowest priority)
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiscordConfig:
|
||||
"""Discord-related configuration."""
|
||||
token: str = ""
|
||||
application_id: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ComfyUIConfig:
|
||||
"""ComfyUI connection configuration."""
|
||||
url: str = "http://127.0.0.1:8188"
|
||||
ws_url: str = "ws://127.0.0.1:8188/ws"
|
||||
timeout: int = 30
|
||||
|
||||
|
||||
@dataclass
|
||||
class DefaultsConfig:
|
||||
"""Default values for bot operations."""
|
||||
max_queue_per_user: int = 3
|
||||
progress_update_interval: float = 2.0
|
||||
workflow_path: Optional[str] = None
|
||||
default_steps: int = 20
|
||||
default_cfg: float = 7.0
|
||||
default_width: int = 512
|
||||
default_height: int = 512
|
||||
|
||||
|
||||
@dataclass
|
||||
class DatabaseConfig:
|
||||
"""Database configuration."""
|
||||
url: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.url:
|
||||
# Default to SQLite in bot/data directory
|
||||
bot_dir = Path(__file__).parent
|
||||
data_dir = bot_dir / "data"
|
||||
data_dir.mkdir(exist_ok=True)
|
||||
self.url = f"sqlite+aiosqlite:///{data_dir}/bot.db"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SecurityConfig:
|
||||
"""Security configuration."""
|
||||
allowed_guilds: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BotConfig:
|
||||
"""Main bot configuration container."""
|
||||
discord: DiscordConfig = field(default_factory=DiscordConfig)
|
||||
comfyui: ComfyUIConfig = field(default_factory=ComfyUIConfig)
|
||||
defaults: DefaultsConfig = field(default_factory=DefaultsConfig)
|
||||
database: DatabaseConfig = field(default_factory=DatabaseConfig)
|
||||
security: SecurityConfig = field(default_factory=SecurityConfig)
|
||||
|
||||
@classmethod
|
||||
def load(cls, config_path: Optional[Path] = None) -> "BotConfig":
|
||||
"""
|
||||
Load configuration from file and environment variables.
|
||||
|
||||
Args:
|
||||
config_path: Optional path to config.yaml file
|
||||
|
||||
Returns:
|
||||
Loaded BotConfig instance
|
||||
"""
|
||||
config = cls()
|
||||
|
||||
# Load from config file if exists
|
||||
if config_path is None:
|
||||
config_path = Path(__file__).parent / "config.yaml"
|
||||
|
||||
if config_path.exists():
|
||||
config._load_from_file(config_path)
|
||||
|
||||
# Override with environment variables
|
||||
config._load_from_env()
|
||||
|
||||
return config
|
||||
|
||||
def _load_from_file(self, path: Path) -> None:
|
||||
"""Load configuration from YAML file."""
|
||||
with open(path) as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
|
||||
# Discord config
|
||||
if "discord" in data:
|
||||
discord_data = data["discord"]
|
||||
if "token" in discord_data:
|
||||
self.discord.token = discord_data["token"]
|
||||
if "application_id" in discord_data:
|
||||
self.discord.application_id = discord_data["application_id"]
|
||||
|
||||
# ComfyUI config
|
||||
if "comfyui" in data:
|
||||
comfyui_data = data["comfyui"]
|
||||
if "url" in comfyui_data:
|
||||
self.comfyui.url = comfyui_data["url"]
|
||||
if "ws_url" in comfyui_data:
|
||||
self.comfyui.ws_url = comfyui_data["ws_url"]
|
||||
if "timeout" in comfyui_data:
|
||||
self.comfyui.timeout = comfyui_data["timeout"]
|
||||
|
||||
# Defaults config
|
||||
if "defaults" in data:
|
||||
defaults_data = data["defaults"]
|
||||
if "max_queue_per_user" in defaults_data:
|
||||
self.defaults.max_queue_per_user = defaults_data["max_queue_per_user"]
|
||||
if "progress_update_interval" in defaults_data:
|
||||
self.defaults.progress_update_interval = defaults_data["progress_update_interval"]
|
||||
if "workflow_path" in defaults_data:
|
||||
self.defaults.workflow_path = defaults_data["workflow_path"]
|
||||
if "default_steps" in defaults_data:
|
||||
self.defaults.default_steps = defaults_data["default_steps"]
|
||||
if "default_cfg" in defaults_data:
|
||||
self.defaults.default_cfg = defaults_data["default_cfg"]
|
||||
if "default_width" in defaults_data:
|
||||
self.defaults.default_width = defaults_data["default_width"]
|
||||
if "default_height" in defaults_data:
|
||||
self.defaults.default_height = defaults_data["default_height"]
|
||||
|
||||
# Database config
|
||||
if "database" in data:
|
||||
db_data = data["database"]
|
||||
if "url" in db_data:
|
||||
self.database.url = db_data["url"]
|
||||
|
||||
# Security config
|
||||
if "security" in data:
|
||||
security_data = data["security"]
|
||||
if "allowed_guilds" in security_data:
|
||||
self.security.allowed_guilds = security_data["allowed_guilds"] or []
|
||||
|
||||
def _load_from_env(self) -> None:
|
||||
"""Load configuration from environment variables."""
|
||||
# Discord
|
||||
if token := os.getenv("DISCORDBOT_DISCORD_TOKEN"):
|
||||
self.discord.token = token
|
||||
if app_id := os.getenv("DISCORDBOT_APPLICATION_ID"):
|
||||
self.discord.application_id = app_id
|
||||
|
||||
# ComfyUI
|
||||
if url := os.getenv("DISCORDBOT_COMFYUI_URL"):
|
||||
self.comfyui.url = url
|
||||
if ws_url := os.getenv("DISCORDBOT_COMFYUI_WS_URL"):
|
||||
self.comfyui.ws_url = ws_url
|
||||
if timeout := os.getenv("DISCORDBOT_COMFYUI_TIMEOUT"):
|
||||
self.comfyui.timeout = int(timeout)
|
||||
|
||||
# Database
|
||||
if db_url := os.getenv("DISCORDBOT_DATABASE_URL"):
|
||||
self.database.url = db_url
|
||||
|
||||
# Defaults
|
||||
if max_queue := os.getenv("DISCORDBOT_MAX_QUEUE_PER_USER"):
|
||||
self.defaults.max_queue_per_user = int(max_queue)
|
||||
if workflow := os.getenv("DISCORDBOT_WORKFLOW_PATH"):
|
||||
self.defaults.workflow_path = workflow
|
||||
|
||||
def validate(self) -> list[str]:
|
||||
"""
|
||||
Validate the configuration.
|
||||
|
||||
Returns:
|
||||
List of validation error messages (empty if valid)
|
||||
"""
|
||||
errors = []
|
||||
|
||||
if not self.discord.token:
|
||||
errors.append("Discord token is required (set DISCORDBOT_DISCORD_TOKEN)")
|
||||
|
||||
if not self.comfyui.url:
|
||||
errors.append("ComfyUI URL is required")
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
# Global config instance (lazy loaded)
|
||||
_config: Optional[BotConfig] = None
|
||||
|
||||
|
||||
def get_config() -> BotConfig:
|
||||
"""Get the global configuration instance."""
|
||||
global _config
|
||||
if _config is None:
|
||||
_config = BotConfig.load()
|
||||
return _config
|
||||
|
||||
|
||||
def reload_config(config_path: Optional[Path] = None) -> BotConfig:
|
||||
"""Reload configuration from disk."""
|
||||
global _config
|
||||
_config = BotConfig.load(config_path)
|
||||
return _config
|
||||
@@ -0,0 +1,107 @@
|
||||
{
|
||||
"3": {
|
||||
"inputs": {
|
||||
"seed": 156680208700286,
|
||||
"steps": 20,
|
||||
"cfg": 8,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1,
|
||||
"model": [
|
||||
"4",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"6",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"7",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"5",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "KSampler",
|
||||
"_meta": {
|
||||
"title": "KSampler"
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"inputs": {
|
||||
"ckpt_name": "v1-5-pruned-emaonly.ckpt"
|
||||
},
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"_meta": {
|
||||
"title": "Load Checkpoint"
|
||||
}
|
||||
},
|
||||
"5": {
|
||||
"inputs": {
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"batch_size": 1
|
||||
},
|
||||
"class_type": "EmptyLatentImage",
|
||||
"_meta": {
|
||||
"title": "Empty Latent Image"
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"inputs": {
|
||||
"text": "beautiful scenery nature glass bottle landscape, , purple galaxy bottle,",
|
||||
"clip": [
|
||||
"4",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "Positive Prompt"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"inputs": {
|
||||
"text": "text, watermark",
|
||||
"clip": [
|
||||
"4",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "Negative Prompt"
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"3",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"4",
|
||||
2
|
||||
]
|
||||
},
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {
|
||||
"title": "VAE Decode"
|
||||
}
|
||||
},
|
||||
"9": {
|
||||
"inputs": {
|
||||
"filename_prefix": "ComfyUI",
|
||||
"images": [
|
||||
"8",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SaveImage",
|
||||
"_meta": {
|
||||
"title": "Save Image"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
"""Database package for bot persistence."""
|
||||
|
||||
from .models import Base, User, Server, ServerRole, Template, Job, Workflow
|
||||
from .repository import Repository
|
||||
|
||||
__all__ = [
|
||||
"Base",
|
||||
"User",
|
||||
"Server",
|
||||
"ServerRole",
|
||||
"Template",
|
||||
"Job",
|
||||
"Workflow",
|
||||
"Repository",
|
||||
]
|
||||
@@ -0,0 +1,222 @@
|
||||
"""
|
||||
SQLAlchemy ORM models for bot persistence.
|
||||
|
||||
Tables:
|
||||
- users: Discord user preferences
|
||||
- servers: Guild configuration
|
||||
- server_roles: Permission role mappings
|
||||
- templates: Saved prompt presets
|
||||
- jobs: Generation history and queue
|
||||
- workflows: Default workflow storage
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import (
|
||||
Boolean,
|
||||
Column,
|
||||
DateTime,
|
||||
Float,
|
||||
ForeignKey,
|
||||
Index,
|
||||
Integer,
|
||||
String,
|
||||
Text,
|
||||
UniqueConstraint,
|
||||
)
|
||||
from sqlalchemy.orm import DeclarativeBase, relationship
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
"""Base class for all models."""
|
||||
pass
|
||||
|
||||
|
||||
class DeliveryType(str, Enum):
|
||||
"""Delivery method for generated images."""
|
||||
CHANNEL = "channel"
|
||||
DM = "dm"
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a generation job."""
|
||||
PENDING = "pending"
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class PermissionLevel(str, Enum):
|
||||
"""Permission levels for role-based access."""
|
||||
USER = "user"
|
||||
GENERATOR = "generator"
|
||||
ADMIN = "admin"
|
||||
|
||||
|
||||
class User(Base):
|
||||
"""Discord user preferences and data."""
|
||||
__tablename__ = "users"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
discord_id = Column(String(20), unique=True, nullable=False, index=True)
|
||||
username = Column(String(100), nullable=False)
|
||||
default_delivery = Column(String(10), default=DeliveryType.CHANNEL.value)
|
||||
default_workflow_id = Column(Integer, ForeignKey("workflows.id"), nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
templates = relationship("Template", back_populates="user", cascade="all, delete-orphan")
|
||||
jobs = relationship("Job", back_populates="user", cascade="all, delete-orphan")
|
||||
default_workflow = relationship("Workflow", foreign_keys=[default_workflow_id])
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<User(id={self.id}, discord_id={self.discord_id}, username={self.username})>"
|
||||
|
||||
|
||||
class Server(Base):
|
||||
"""Discord server (guild) configuration."""
|
||||
__tablename__ = "servers"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
discord_id = Column(String(20), unique=True, nullable=False, index=True)
|
||||
name = Column(String(100), nullable=False)
|
||||
default_channel_id = Column(String(20), nullable=True)
|
||||
enabled = Column(Boolean, default=True)
|
||||
max_queue_per_user = Column(Integer, default=3)
|
||||
default_workflow_id = Column(Integer, ForeignKey("workflows.id"), nullable=True)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
roles = relationship("ServerRole", back_populates="server", cascade="all, delete-orphan")
|
||||
templates = relationship("Template", back_populates="server", cascade="all, delete-orphan")
|
||||
jobs = relationship("Job", back_populates="server")
|
||||
default_workflow = relationship("Workflow", foreign_keys=[default_workflow_id])
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<Server(id={self.id}, discord_id={self.discord_id}, name={self.name})>"
|
||||
|
||||
|
||||
class ServerRole(Base):
|
||||
"""Role-based permission mapping for servers."""
|
||||
__tablename__ = "server_roles"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
server_id = Column(Integer, ForeignKey("servers.id"), nullable=False)
|
||||
role_discord_id = Column(String(20), nullable=False)
|
||||
permission_level = Column(String(20), nullable=False)
|
||||
|
||||
# Relationships
|
||||
server = relationship("Server", back_populates="roles")
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("server_id", "role_discord_id", name="uq_server_role"),
|
||||
Index("idx_server_roles_server", "server_id"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<ServerRole(server_id={self.server_id}, role={self.role_discord_id}, level={self.permission_level})>"
|
||||
|
||||
|
||||
class Workflow(Base):
|
||||
"""Stored workflow configurations."""
|
||||
__tablename__ = "workflows"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
name = Column(String(100), nullable=False)
|
||||
description = Column(Text, nullable=True)
|
||||
workflow_json = Column(Text, nullable=False)
|
||||
is_default = Column(Boolean, default=False)
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<Workflow(id={self.id}, name={self.name}, is_default={self.is_default})>"
|
||||
|
||||
|
||||
class Template(Base):
|
||||
"""User-saved prompt templates."""
|
||||
__tablename__ = "templates"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
server_id = Column(Integer, ForeignKey("servers.id"), nullable=True) # NULL = private
|
||||
name = Column(String(100), nullable=False)
|
||||
positive_prompt = Column(Text, nullable=False)
|
||||
negative_prompt = Column(Text, default="")
|
||||
parameters = Column(Text, nullable=True) # JSON: {steps, cfg, seed, etc.}
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="templates")
|
||||
server = relationship("Server", back_populates="templates")
|
||||
|
||||
__table_args__ = (
|
||||
Index("idx_templates_user", "user_id"),
|
||||
Index("idx_templates_server", "server_id"),
|
||||
UniqueConstraint("user_id", "server_id", "name", name="uq_template_name"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<Template(id={self.id}, name={self.name}, user_id={self.user_id})>"
|
||||
|
||||
|
||||
class Job(Base):
|
||||
"""Generation job tracking and history."""
|
||||
__tablename__ = "jobs"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
prompt_id = Column(String(64), unique=True, nullable=False, index=True)
|
||||
user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
|
||||
server_id = Column(Integer, ForeignKey("servers.id"), nullable=True)
|
||||
channel_id = Column(String(20), nullable=True)
|
||||
|
||||
# Status tracking
|
||||
status = Column(String(20), default=JobStatus.PENDING.value, index=True)
|
||||
queue_position = Column(Integer, nullable=True)
|
||||
progress = Column(Integer, default=0)
|
||||
progress_max = Column(Integer, default=0)
|
||||
|
||||
# Generation parameters
|
||||
positive_prompt = Column(Text, nullable=True)
|
||||
negative_prompt = Column(Text, nullable=True)
|
||||
parameters = Column(Text, nullable=True) # JSON
|
||||
workflow_json = Column(Text, nullable=True)
|
||||
|
||||
# Results
|
||||
output_images = Column(Text, nullable=True) # JSON array
|
||||
error_message = Column(Text, nullable=True)
|
||||
|
||||
# Delivery
|
||||
delivery_type = Column(String(10), default=DeliveryType.CHANNEL.value)
|
||||
message_id = Column(String(20), nullable=True) # Discord message for updates
|
||||
|
||||
# Timestamps
|
||||
created_at = Column(DateTime, default=datetime.utcnow)
|
||||
started_at = Column(DateTime, nullable=True)
|
||||
completed_at = Column(DateTime, nullable=True)
|
||||
|
||||
# Relationships
|
||||
user = relationship("User", back_populates="jobs")
|
||||
server = relationship("Server", back_populates="jobs")
|
||||
|
||||
__table_args__ = (
|
||||
Index("idx_jobs_user", "user_id"),
|
||||
Index("idx_jobs_status", "status"),
|
||||
Index("idx_jobs_created", "created_at"),
|
||||
)
|
||||
|
||||
@property
|
||||
def duration(self) -> Optional[float]:
|
||||
"""Get job duration in seconds."""
|
||||
if self.started_at and self.completed_at:
|
||||
return (self.completed_at - self.started_at).total_seconds()
|
||||
return None
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<Job(id={self.id}, prompt_id={self.prompt_id}, status={self.status})>"
|
||||
@@ -0,0 +1,656 @@
|
||||
"""
|
||||
Data access layer for bot database operations.
|
||||
|
||||
Provides async CRUD operations for all database models.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import select, update, delete, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
|
||||
from sqlalchemy.orm import selectinload
|
||||
|
||||
from .models import (
|
||||
Base,
|
||||
User,
|
||||
Server,
|
||||
ServerRole,
|
||||
Template,
|
||||
Job,
|
||||
Workflow,
|
||||
JobStatus,
|
||||
PermissionLevel,
|
||||
)
|
||||
|
||||
|
||||
class Repository:
|
||||
"""Async repository for database operations."""
|
||||
|
||||
def __init__(self, database_url: str):
|
||||
"""
|
||||
Initialize the repository.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy async database URL
|
||||
"""
|
||||
self.engine = create_async_engine(database_url, echo=False)
|
||||
self.async_session = async_sessionmaker(
|
||||
self.engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
async def init_db(self) -> None:
|
||||
"""Create all database tables."""
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the database connection."""
|
||||
await self.engine.dispose()
|
||||
|
||||
# ==================== User Operations ====================
|
||||
|
||||
async def get_or_create_user(
|
||||
self,
|
||||
discord_id: str,
|
||||
username: str,
|
||||
) -> User:
|
||||
"""Get existing user or create new one."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(User).where(User.discord_id == discord_id)
|
||||
)
|
||||
user = result.scalar_one_or_none()
|
||||
|
||||
if user is None:
|
||||
user = User(discord_id=discord_id, username=username)
|
||||
session.add(user)
|
||||
await session.commit()
|
||||
await session.refresh(user)
|
||||
elif user.username != username and username != "Unknown":
|
||||
user.username = username
|
||||
await session.commit()
|
||||
|
||||
return user
|
||||
|
||||
async def get_user(self, discord_id: str) -> Optional[User]:
|
||||
"""Get user by Discord ID."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(User).where(User.discord_id == discord_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def update_user_delivery(
|
||||
self,
|
||||
discord_id: str,
|
||||
delivery_type: str,
|
||||
) -> Optional[User]:
|
||||
"""Update user's default delivery preference."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(User).where(User.discord_id == discord_id)
|
||||
)
|
||||
user = result.scalar_one_or_none()
|
||||
if user:
|
||||
user.default_delivery = delivery_type
|
||||
await session.commit()
|
||||
return user
|
||||
|
||||
# ==================== Server Operations ====================
|
||||
|
||||
async def get_or_create_server(
|
||||
self,
|
||||
discord_id: str,
|
||||
name: str,
|
||||
) -> Server:
|
||||
"""Get existing server or create new one."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Server).where(Server.discord_id == discord_id)
|
||||
)
|
||||
server = result.scalar_one_or_none()
|
||||
|
||||
if server is None:
|
||||
server = Server(discord_id=discord_id, name=name)
|
||||
session.add(server)
|
||||
await session.commit()
|
||||
await session.refresh(server)
|
||||
elif server.name != name:
|
||||
server.name = name
|
||||
await session.commit()
|
||||
|
||||
return server
|
||||
|
||||
async def get_server(self, discord_id: str) -> Optional[Server]:
|
||||
"""Get server by Discord ID."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Server)
|
||||
.options(selectinload(Server.roles))
|
||||
.where(Server.discord_id == discord_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def update_server_channel(
|
||||
self,
|
||||
discord_id: str,
|
||||
channel_id: str,
|
||||
) -> Optional[Server]:
|
||||
"""Update server's default output channel."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Server).where(Server.discord_id == discord_id)
|
||||
)
|
||||
server = result.scalar_one_or_none()
|
||||
if server:
|
||||
server.default_channel_id = channel_id
|
||||
await session.commit()
|
||||
return server
|
||||
|
||||
async def update_server_queue_limit(
|
||||
self,
|
||||
discord_id: str,
|
||||
limit: int,
|
||||
) -> Optional[Server]:
|
||||
"""Update server's per-user queue limit."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Server).where(Server.discord_id == discord_id)
|
||||
)
|
||||
server = result.scalar_one_or_none()
|
||||
if server:
|
||||
server.max_queue_per_user = limit
|
||||
await session.commit()
|
||||
return server
|
||||
|
||||
# ==================== Role Operations ====================
|
||||
|
||||
async def set_server_role(
|
||||
self,
|
||||
server_discord_id: str,
|
||||
role_discord_id: str,
|
||||
permission_level: str,
|
||||
) -> ServerRole:
|
||||
"""Set or update a role's permission level for a server."""
|
||||
async with self.async_session() as session:
|
||||
# Get server
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if not server:
|
||||
raise ValueError(f"Server {server_discord_id} not found")
|
||||
|
||||
# Check for existing role mapping
|
||||
role_result = await session.execute(
|
||||
select(ServerRole).where(
|
||||
ServerRole.server_id == server.id,
|
||||
ServerRole.role_discord_id == role_discord_id,
|
||||
)
|
||||
)
|
||||
role = role_result.scalar_one_or_none()
|
||||
|
||||
if role:
|
||||
role.permission_level = permission_level
|
||||
else:
|
||||
role = ServerRole(
|
||||
server_id=server.id,
|
||||
role_discord_id=role_discord_id,
|
||||
permission_level=permission_level,
|
||||
)
|
||||
session.add(role)
|
||||
|
||||
await session.commit()
|
||||
await session.refresh(role)
|
||||
return role
|
||||
|
||||
async def get_server_roles(self, server_discord_id: str) -> list[ServerRole]:
|
||||
"""Get all role mappings for a server."""
|
||||
async with self.async_session() as session:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if not server:
|
||||
return []
|
||||
|
||||
result = await session.execute(
|
||||
select(ServerRole).where(ServerRole.server_id == server.id)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def delete_server_role(
|
||||
self,
|
||||
server_discord_id: str,
|
||||
role_discord_id: str,
|
||||
) -> bool:
|
||||
"""Remove a role mapping."""
|
||||
async with self.async_session() as session:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if not server:
|
||||
return False
|
||||
|
||||
result = await session.execute(
|
||||
delete(ServerRole).where(
|
||||
ServerRole.server_id == server.id,
|
||||
ServerRole.role_discord_id == role_discord_id,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
return result.rowcount > 0
|
||||
|
||||
# ==================== Template Operations ====================
|
||||
|
||||
async def create_template(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
name: str,
|
||||
positive_prompt: str,
|
||||
negative_prompt: str = "",
|
||||
parameters: Optional[dict] = None,
|
||||
server_discord_id: Optional[str] = None,
|
||||
) -> Template:
|
||||
"""Create a new prompt template."""
|
||||
async with self.async_session() as session:
|
||||
# Get user
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
raise ValueError(f"User {user_discord_id} not found")
|
||||
|
||||
# Get server if provided
|
||||
server_id = None
|
||||
if server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
server_id = server.id
|
||||
|
||||
template = Template(
|
||||
user_id=user.id,
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
positive_prompt=positive_prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
parameters=json.dumps(parameters) if parameters else None,
|
||||
)
|
||||
session.add(template)
|
||||
await session.commit()
|
||||
await session.refresh(template)
|
||||
return template
|
||||
|
||||
async def get_template(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
name: str,
|
||||
server_discord_id: Optional[str] = None,
|
||||
) -> Optional[Template]:
|
||||
"""Get a template by name."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
return None
|
||||
|
||||
# Build query
|
||||
query = select(Template).where(
|
||||
Template.user_id == user.id,
|
||||
Template.name == name,
|
||||
)
|
||||
|
||||
if server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
query = query.where(Template.server_id == server.id)
|
||||
else:
|
||||
query = query.where(Template.server_id.is_(None))
|
||||
|
||||
result = await session.execute(query)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_templates(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
server_discord_id: Optional[str] = None,
|
||||
include_shared: bool = True,
|
||||
) -> list[Template]:
|
||||
"""List templates for a user."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
return []
|
||||
|
||||
# Private templates
|
||||
query = select(Template).where(
|
||||
Template.user_id == user.id,
|
||||
Template.server_id.is_(None),
|
||||
)
|
||||
result = await session.execute(query)
|
||||
templates = list(result.scalars().all())
|
||||
|
||||
# Shared templates in server
|
||||
if include_shared and server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
shared_query = select(Template).where(
|
||||
Template.server_id == server.id
|
||||
)
|
||||
shared_result = await session.execute(shared_query)
|
||||
templates.extend(shared_result.scalars().all())
|
||||
|
||||
return templates
|
||||
|
||||
async def delete_template(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
name: str,
|
||||
server_discord_id: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Delete a template."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
return False
|
||||
|
||||
query = delete(Template).where(
|
||||
Template.user_id == user.id,
|
||||
Template.name == name,
|
||||
)
|
||||
|
||||
if server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
query = query.where(Template.server_id == server.id)
|
||||
else:
|
||||
query = query.where(Template.server_id.is_(None))
|
||||
|
||||
result = await session.execute(query)
|
||||
await session.commit()
|
||||
return result.rowcount > 0
|
||||
|
||||
# ==================== Job Operations ====================
|
||||
|
||||
async def create_job(
|
||||
self,
|
||||
prompt_id: str,
|
||||
user_discord_id: str,
|
||||
positive_prompt: str,
|
||||
negative_prompt: str = "",
|
||||
parameters: Optional[dict] = None,
|
||||
workflow_json: Optional[str] = None,
|
||||
delivery_type: str = "channel",
|
||||
server_discord_id: Optional[str] = None,
|
||||
channel_id: Optional[str] = None,
|
||||
) -> Job:
|
||||
"""Create a new generation job."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
raise ValueError(f"User {user_discord_id} not found")
|
||||
|
||||
server_id = None
|
||||
if server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
server_id = server.id
|
||||
|
||||
job = Job(
|
||||
prompt_id=prompt_id,
|
||||
user_id=user.id,
|
||||
server_id=server_id,
|
||||
channel_id=channel_id,
|
||||
positive_prompt=positive_prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
parameters=json.dumps(parameters) if parameters else None,
|
||||
workflow_json=workflow_json,
|
||||
delivery_type=delivery_type,
|
||||
)
|
||||
session.add(job)
|
||||
await session.commit()
|
||||
await session.refresh(job)
|
||||
return job
|
||||
|
||||
async def get_job(self, prompt_id: str) -> Optional[Job]:
|
||||
"""Get a job by prompt ID."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job)
|
||||
.options(selectinload(Job.user))
|
||||
.where(Job.prompt_id == prompt_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_job_by_id(self, job_id: int) -> Optional[Job]:
|
||||
"""Get a job by internal ID."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job)
|
||||
.options(selectinload(Job.user))
|
||||
.where(Job.id == job_id)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def update_job_status(
|
||||
self,
|
||||
prompt_id: str,
|
||||
status: str,
|
||||
error_message: Optional[str] = None,
|
||||
output_images: Optional[list[str]] = None,
|
||||
) -> Optional[Job]:
|
||||
"""Update job status."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job)
|
||||
.options(selectinload(Job.user))
|
||||
.where(Job.prompt_id == prompt_id)
|
||||
)
|
||||
job = result.scalar_one_or_none()
|
||||
if not job:
|
||||
return None
|
||||
|
||||
job.status = status
|
||||
|
||||
if status == JobStatus.RUNNING.value:
|
||||
job.started_at = datetime.utcnow()
|
||||
elif status in (JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value):
|
||||
job.completed_at = datetime.utcnow()
|
||||
|
||||
if error_message is not None:
|
||||
job.error_message = error_message
|
||||
|
||||
if output_images is not None:
|
||||
current_images = []
|
||||
if job.output_images:
|
||||
try:
|
||||
current_images = json.loads(job.output_images)
|
||||
except json.JSONDecodeError:
|
||||
current_images = []
|
||||
|
||||
# Append new images
|
||||
current_images.extend(output_images)
|
||||
job.output_images = json.dumps(current_images)
|
||||
|
||||
await session.commit()
|
||||
return job
|
||||
|
||||
async def update_job_progress(
|
||||
self,
|
||||
prompt_id: str,
|
||||
progress: int,
|
||||
progress_max: int,
|
||||
) -> Optional[Job]:
|
||||
"""Update job progress."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job).where(Job.prompt_id == prompt_id)
|
||||
)
|
||||
job = result.scalar_one_or_none()
|
||||
if job:
|
||||
job.progress = progress
|
||||
job.progress_max = progress_max
|
||||
await session.commit()
|
||||
return job
|
||||
|
||||
async def update_job_message(
|
||||
self,
|
||||
prompt_id: str,
|
||||
message_id: str,
|
||||
) -> Optional[Job]:
|
||||
"""Update the Discord message ID for a job."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job).where(Job.prompt_id == prompt_id)
|
||||
)
|
||||
job = result.scalar_one_or_none()
|
||||
if job:
|
||||
job.message_id = message_id
|
||||
await session.commit()
|
||||
return job
|
||||
|
||||
async def list_user_jobs(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
limit: int = 10,
|
||||
status: Optional[str] = None,
|
||||
) -> list[Job]:
|
||||
"""List jobs for a user."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
return []
|
||||
|
||||
query = (
|
||||
select(Job)
|
||||
.where(Job.user_id == user.id)
|
||||
.order_by(Job.created_at.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
|
||||
if status:
|
||||
query = query.where(Job.status == status)
|
||||
|
||||
result = await session.execute(query)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def count_user_pending_jobs(
|
||||
self,
|
||||
user_discord_id: str,
|
||||
server_discord_id: Optional[str] = None,
|
||||
) -> int:
|
||||
"""Count pending/running jobs for a user."""
|
||||
async with self.async_session() as session:
|
||||
user_result = await session.execute(
|
||||
select(User).where(User.discord_id == user_discord_id)
|
||||
)
|
||||
user = user_result.scalar_one_or_none()
|
||||
if not user:
|
||||
return 0
|
||||
|
||||
query = (
|
||||
select(func.count(Job.id))
|
||||
.where(Job.user_id == user.id)
|
||||
.where(Job.status.in_([JobStatus.PENDING.value, JobStatus.RUNNING.value]))
|
||||
)
|
||||
|
||||
if server_discord_id:
|
||||
server_result = await session.execute(
|
||||
select(Server).where(Server.discord_id == server_discord_id)
|
||||
)
|
||||
server = server_result.scalar_one_or_none()
|
||||
if server:
|
||||
query = query.where(Job.server_id == server.id)
|
||||
|
||||
result = await session.execute(query)
|
||||
return result.scalar() or 0
|
||||
|
||||
async def get_pending_jobs(self) -> list[Job]:
|
||||
"""Get all pending jobs ordered by creation time."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Job)
|
||||
.options(selectinload(Job.user))
|
||||
.where(Job.status.in_([JobStatus.PENDING.value, JobStatus.RUNNING.value]))
|
||||
.order_by(Job.created_at)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
# ==================== Workflow Operations ====================
|
||||
|
||||
async def save_workflow(
|
||||
self,
|
||||
name: str,
|
||||
workflow_json: str,
|
||||
description: Optional[str] = None,
|
||||
is_default: bool = False,
|
||||
) -> Workflow:
|
||||
"""Save a workflow configuration."""
|
||||
async with self.async_session() as session:
|
||||
# If setting as default, unset other defaults
|
||||
if is_default:
|
||||
await session.execute(
|
||||
update(Workflow).where(Workflow.is_default == True).values(is_default=False)
|
||||
)
|
||||
|
||||
workflow = Workflow(
|
||||
name=name,
|
||||
workflow_json=workflow_json,
|
||||
description=description,
|
||||
is_default=is_default,
|
||||
)
|
||||
session.add(workflow)
|
||||
await session.commit()
|
||||
await session.refresh(workflow)
|
||||
return workflow
|
||||
|
||||
async def get_default_workflow(self) -> Optional[Workflow]:
|
||||
"""Get the default workflow."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Workflow).where(Workflow.is_default == True)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_workflow(self, name: str) -> Optional[Workflow]:
|
||||
"""Get a workflow by name."""
|
||||
async with self.async_session() as session:
|
||||
result = await session.execute(
|
||||
select(Workflow).where(Workflow.name == name)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Discord embed builders."""
|
||||
|
||||
from .builders import EmbedBuilder
|
||||
|
||||
__all__ = ["EmbedBuilder"]
|
||||
@@ -0,0 +1,72 @@
|
||||
import discord
|
||||
import logging
|
||||
from typing import Optional, List
|
||||
from datetime import datetime
|
||||
|
||||
class EmbedBuilder:
|
||||
"""Helper for building Discord embeds."""
|
||||
|
||||
@staticmethod
|
||||
def job_queued(job, position: int = 0) -> discord.Embed:
|
||||
"""Embed for queued job."""
|
||||
embed = discord.Embed(
|
||||
title="🎨 Generation Queued",
|
||||
description=f"**Prompt:** {job.positive_prompt}",
|
||||
color=discord.Color.blue(),
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
embed.add_field(name="Queue Position", value=str(position) if position > 0 else "Pending...", inline=True)
|
||||
embed.add_field(name="Status", value="Waiting to start...", inline=True)
|
||||
if job.negative_prompt:
|
||||
embed.add_field(name="Negative Prompt", value=job.negative_prompt, inline=False)
|
||||
embed.set_footer(text=f"Job ID: {job.id}")
|
||||
return embed
|
||||
|
||||
@staticmethod
|
||||
def job_progress(job, progress: int, max_progress: int) -> discord.Embed:
|
||||
"""Embed for running job with progress."""
|
||||
percent = int((progress / max_progress) * 100) if max_progress > 0 else 0
|
||||
bars = "█" * (percent // 10) + "░" * (10 - (percent // 10))
|
||||
|
||||
embed = discord.Embed(
|
||||
title="🎨 Generating...",
|
||||
description=f"**Prompt:** {job.positive_prompt}",
|
||||
color=discord.Color.orange(),
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
embed.add_field(name="Progress", value=f"`{bars}` {percent}%", inline=False)
|
||||
if job.negative_prompt:
|
||||
embed.add_field(name="Negative Prompt", value=job.negative_prompt, inline=False)
|
||||
embed.set_footer(text=f"Job ID: {job.id}")
|
||||
return embed
|
||||
|
||||
@staticmethod
|
||||
def job_completed(job, image_count: int) -> discord.Embed:
|
||||
"""Embed for completed job."""
|
||||
embed = discord.Embed(
|
||||
title="✨ Generation Complete!",
|
||||
description=f"**Prompt:** {job.positive_prompt}",
|
||||
color=discord.Color.green(),
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
embed.add_field(name="Images", value=f"{image_count} generated", inline=True)
|
||||
embed.add_field(name="Duration", value=f"{job.duration:.1f}s" if hasattr(job, 'duration') and job.duration else "Done", inline=True)
|
||||
|
||||
if job.negative_prompt:
|
||||
embed.add_field(name="Negative Prompt", value=job.negative_prompt, inline=False)
|
||||
|
||||
embed.set_footer(text=f"Job ID: {job.id}")
|
||||
return embed
|
||||
|
||||
@staticmethod
|
||||
def job_failed(job, error_message: str) -> discord.Embed:
|
||||
"""Embed for failed job."""
|
||||
embed = discord.Embed(
|
||||
title="❌ Generation Failed",
|
||||
description=f"**Prompt:** {job.positive_prompt}",
|
||||
color=discord.Color.red(),
|
||||
timestamp=datetime.utcnow()
|
||||
)
|
||||
embed.add_field(name="Error", value=f"```{error_message}```", inline=False)
|
||||
embed.set_footer(text=f"Job ID: {job.id}")
|
||||
return embed
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Bot services for business logic."""
|
||||
|
||||
from .permissions import PermissionService, Permissions, require_permission
|
||||
from .job_manager import JobManager
|
||||
from .delivery import DeliveryService
|
||||
|
||||
__all__ = [
|
||||
"PermissionService",
|
||||
"Permissions",
|
||||
"require_permission",
|
||||
"JobManager",
|
||||
"DeliveryService",
|
||||
]
|
||||
@@ -0,0 +1,120 @@
|
||||
import discord
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from ..comfyui.client import ComfyUIClient
|
||||
from ..embeds.builders import EmbedBuilder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DeliveryService:
|
||||
"""Handles delivery of results to Discord."""
|
||||
|
||||
def __init__(self, bot: discord.Client, comfy_client: ComfyUIClient):
|
||||
self.bot = bot
|
||||
self.client = comfy_client
|
||||
|
||||
async def _get_destination(
|
||||
self, job: Any
|
||||
) -> Optional[Union[discord.User, discord.TextChannel]]:
|
||||
"""Get the destination channel or user for a job."""
|
||||
if job.delivery_type == "dm":
|
||||
try:
|
||||
user = self.bot.get_user(int(job.user.discord_id))
|
||||
if not user:
|
||||
user = await self.bot.fetch_user(int(job.user.discord_id))
|
||||
return user
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch user for DM: {e}")
|
||||
return None
|
||||
|
||||
if job.channel_id:
|
||||
destination = self.bot.get_channel(int(job.channel_id))
|
||||
if destination:
|
||||
return destination
|
||||
# Fallback to DM if channel not found
|
||||
try:
|
||||
user = self.bot.get_user(int(job.user.discord_id))
|
||||
if not user:
|
||||
user = await self.bot.fetch_user(int(job.user.discord_id))
|
||||
return user
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
async def deliver_error(self, job: Any, error_message: str) -> bool:
|
||||
"""
|
||||
Deliver error notification for a failed job.
|
||||
|
||||
Args:
|
||||
job: The failed Job model instance
|
||||
error_message: The error message to display
|
||||
|
||||
Returns:
|
||||
True if delivery succeeded, False otherwise
|
||||
"""
|
||||
# Truncate error message if too long (Discord embed field limit)
|
||||
if len(error_message) > 1000:
|
||||
error_message = error_message[:997] + "..."
|
||||
|
||||
embed = EmbedBuilder.job_failed(job, error_message)
|
||||
|
||||
destination = await self._get_destination(job)
|
||||
|
||||
if destination:
|
||||
try:
|
||||
await destination.send(embed=embed)
|
||||
logger.info(f"Delivered error for job {job.id} to {destination}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deliver error: {e}")
|
||||
return False
|
||||
else:
|
||||
logger.error(
|
||||
f"Could not determine destination for error delivery (job {job.id})"
|
||||
)
|
||||
return False
|
||||
|
||||
async def deliver_job(self, job: Any): # job: Job model
|
||||
"""Deliver results for a completed job."""
|
||||
if not job.output_images:
|
||||
logger.warning(f"Job {job.id} completed but has no images.")
|
||||
return
|
||||
|
||||
try:
|
||||
images_data = json.loads(job.output_images)
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to parse output images for job {job.id}")
|
||||
return
|
||||
|
||||
files = []
|
||||
for img_meta in images_data:
|
||||
filename = img_meta.get("filename")
|
||||
subfolder = img_meta.get("subfolder", "")
|
||||
img_type = img_meta.get("type", "output")
|
||||
|
||||
try:
|
||||
img_bytes = await self.client.get_image(filename, subfolder, img_type)
|
||||
files.append(discord.File(io.BytesIO(img_bytes), filename=filename))
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to download image {filename}: {e}")
|
||||
|
||||
if not files:
|
||||
logger.warning("No files to upload.")
|
||||
return
|
||||
|
||||
destination = await self._get_destination(job)
|
||||
|
||||
if destination:
|
||||
content = f"Generation complete for <@{job.user.discord_id}>!\n**Prompt:** {job.positive_prompt}"
|
||||
try:
|
||||
await destination.send(content=content, files=files)
|
||||
logger.info(f"Delivered job {job.id} to {destination}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send message: {e}")
|
||||
else:
|
||||
logger.error(f"Could not determine destination for job {job.id}")
|
||||
@@ -0,0 +1,208 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Optional, Dict, List, Any
|
||||
|
||||
from ..database.repository import Repository
|
||||
from ..database.models import JobStatus, Job
|
||||
from ..comfyui.client import ComfyUIClient
|
||||
from ..comfyui.websocket import ComfyUIWebSocket
|
||||
from .delivery import DeliveryService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class JobManager:
|
||||
"""Manages the lifecycle of generation jobs."""
|
||||
|
||||
def __init__(self, repository: Repository, comfy_client: ComfyUIClient, comfy_ws: ComfyUIWebSocket, delivery_service: DeliveryService):
|
||||
self.repo = repository
|
||||
self.client = comfy_client
|
||||
self.ws = comfy_ws
|
||||
self.delivery = delivery_service
|
||||
# Use common client_id from WebSocket
|
||||
self.client_id = comfy_ws.client_id
|
||||
|
||||
# In-memory mapping of prompt_id -> current status buffer
|
||||
self._active_jobs = {}
|
||||
self._delivery_tasks = {} # prompt_id -> asyncio.Task
|
||||
|
||||
async def start(self):
|
||||
"""Start listening to WebSocket events."""
|
||||
self.ws.add_listener("status", self._on_status)
|
||||
self.ws.add_listener("execution_start", self._on_execution_start)
|
||||
self.ws.add_listener("executing", self._on_executing)
|
||||
self.ws.add_listener("executed", self._on_executed)
|
||||
self.ws.add_listener("execution_error", self._on_execution_error)
|
||||
self.ws.add_listener("progress", self._on_progress)
|
||||
logger.info(f"JobManager started with client_id: {self.client_id}")
|
||||
|
||||
async def create_job(self,
|
||||
user_discord_id: str,
|
||||
workflow: Dict[str, Any],
|
||||
positive_prompt: str,
|
||||
negative_prompt: str = "",
|
||||
parameters: Optional[Dict] = None,
|
||||
server_discord_id: Optional[str] = None,
|
||||
channel_id: Optional[str] = None,
|
||||
delivery_type: str = "channel") -> Job:
|
||||
"""Submit a job to ComfyUI and database."""
|
||||
|
||||
# 0. Validate and ensure entities exist
|
||||
user = await self.repo.get_or_create_user(user_discord_id, "Unknown")
|
||||
|
||||
# Check queue limit
|
||||
current_queue_count = await self.repo.count_user_pending_jobs(user_discord_id, server_discord_id)
|
||||
|
||||
# Get server specific limit or default
|
||||
max_queue = 3
|
||||
if server_discord_id:
|
||||
server = await self.repo.get_server(server_discord_id)
|
||||
if server:
|
||||
max_queue = server.max_queue_per_user
|
||||
|
||||
if current_queue_count >= max_queue:
|
||||
raise ValueError(f"Queue limit reached ({max_queue} jobs). Please wait for your current jobs to finish.")
|
||||
|
||||
# 1. Submit to ComfyUI
|
||||
response = await self.client.queue_prompt(workflow, self.client_id)
|
||||
prompt_id = response.get("prompt_id")
|
||||
|
||||
if not prompt_id:
|
||||
raise ValueError("Failed to get prompt_id from ComfyUI")
|
||||
|
||||
# 2. Create DB entry
|
||||
job = await self.repo.create_job(
|
||||
prompt_id=prompt_id,
|
||||
user_discord_id=user_discord_id,
|
||||
server_discord_id=server_discord_id,
|
||||
channel_id=channel_id,
|
||||
positive_prompt=positive_prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
parameters=parameters,
|
||||
workflow_json=json.dumps(workflow),
|
||||
delivery_type=delivery_type
|
||||
)
|
||||
|
||||
logger.info(f"Created job {job.id} (prompt_id: {prompt_id}) for user {user_discord_id}")
|
||||
return job
|
||||
|
||||
async def cancel_job(self, job_id: int) -> bool:
|
||||
"""Cancel a job."""
|
||||
job = await self.repo.get_job_by_id(job_id)
|
||||
if not job:
|
||||
return False
|
||||
|
||||
if job.status in [JobStatus.COMPLETED.value, JobStatus.FAILED.value, JobStatus.CANCELLED.value]:
|
||||
return False
|
||||
|
||||
if job.status == JobStatus.RUNNING.value:
|
||||
await self.client.interrupt()
|
||||
|
||||
# Remove from queue if pending
|
||||
try:
|
||||
await self.client.delete_queue_item(job.prompt_id)
|
||||
except:
|
||||
pass
|
||||
|
||||
await self.repo.update_job_status(job.prompt_id, JobStatus.CANCELLED.value)
|
||||
return True
|
||||
|
||||
async def _schedule_delivery(self, job: Job):
|
||||
"""Schedule a debounced delivery for a job."""
|
||||
prompt_id = job.prompt_id
|
||||
|
||||
# Cancel existing timer if any
|
||||
if prompt_id in self._delivery_tasks:
|
||||
self._delivery_tasks[prompt_id].cancel()
|
||||
|
||||
# Create new timer
|
||||
task = asyncio.create_task(self._deliver_delayed(job))
|
||||
self._delivery_tasks[prompt_id] = task
|
||||
|
||||
# Cleanup callback
|
||||
def cleanup(t):
|
||||
if prompt_id in self._delivery_tasks and self._delivery_tasks[prompt_id] == t:
|
||||
self._delivery_tasks.pop(prompt_id, None)
|
||||
|
||||
task.add_done_callback(cleanup)
|
||||
|
||||
async def _deliver_delayed(self, job: Job):
|
||||
"""Wait briefly then deliver the job."""
|
||||
try:
|
||||
await asyncio.sleep(1.0) # Debounce window
|
||||
logger.info(f"Job {job.id} delivery timer expired. Delivering results...")
|
||||
await self.delivery.deliver_job(job)
|
||||
except asyncio.CancelledError:
|
||||
logger.debug(f"Job {job.id} delivery debounced/cancelled.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in delayed delivery for job {job.id}: {e}")
|
||||
|
||||
# -- WebSocket Event Handlers --
|
||||
|
||||
async def _on_status(self, data: Dict[str, Any]):
|
||||
pass
|
||||
|
||||
async def _on_execution_start(self, data: Dict[str, Any]):
|
||||
prompt_id = data.get("data", {}).get("prompt_id")
|
||||
if prompt_id:
|
||||
await self.repo.update_job_status(prompt_id, JobStatus.RUNNING.value)
|
||||
|
||||
async def _on_executing(self, data: Dict[str, Any]):
|
||||
pass
|
||||
|
||||
async def _on_progress(self, data: Dict[str, Any]):
|
||||
msg = data.get("data", {})
|
||||
prompt_id = msg.get("prompt_id")
|
||||
value = msg.get("value")
|
||||
max_val = msg.get("max")
|
||||
|
||||
if prompt_id and value is not None and max_val is not None:
|
||||
await self.repo.update_job_progress(prompt_id, value, max_val)
|
||||
|
||||
async def _on_executed(self, data: Dict[str, Any]):
|
||||
"""
|
||||
Handle execution completion of a node.
|
||||
If it contains images, we assume it's a relevant output.
|
||||
We accumulate images and update job status.
|
||||
"""
|
||||
msg = data.get("data", {})
|
||||
prompt_id = msg.get("prompt_id")
|
||||
output = msg.get("output", {})
|
||||
|
||||
if prompt_id and "images" in output:
|
||||
# Found images
|
||||
images = output["images"]
|
||||
|
||||
# Update DB with images and mark as completed
|
||||
# NOTE: In complex workflows, there might be multiple outputs.
|
||||
# Ideally we check if this is the last one or something.
|
||||
# But normally 'executed' with images means we got something.
|
||||
# We'll mark as completed for now. If multiple exist, latest wins.
|
||||
|
||||
job = await self.repo.update_job_status(
|
||||
prompt_id,
|
||||
JobStatus.COMPLETED.value,
|
||||
output_images=images
|
||||
)
|
||||
|
||||
if job:
|
||||
logger.info(f"Job {job.id} update received. Scheduling delivery...")
|
||||
await self._schedule_delivery(job)
|
||||
|
||||
async def _on_execution_error(self, data: Dict[str, Any]):
|
||||
msg = data.get("data", {})
|
||||
prompt_id = msg.get("prompt_id")
|
||||
exception_type = msg.get("exception_type", "Unknown Error")
|
||||
exception_message = msg.get("exception_message", "")
|
||||
|
||||
if prompt_id:
|
||||
error_msg = f"{exception_type}: {exception_message}"
|
||||
job = await self.repo.update_job_status(
|
||||
prompt_id, JobStatus.FAILED.value, error_message=error_msg
|
||||
)
|
||||
|
||||
# Deliver error notification to user
|
||||
if job:
|
||||
await self.delivery.deliver_error(job, error_msg)
|
||||
@@ -0,0 +1,95 @@
|
||||
import logging
|
||||
from typing import Optional, List, Dict, Union
|
||||
import discord
|
||||
from enum import Enum
|
||||
|
||||
from ..database.repository import Repository
|
||||
from ..database.models import PermissionLevel, ServerRole, User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class Permissions(Enum):
|
||||
USER = "user"
|
||||
GENERATOR = "generator"
|
||||
ADMIN = "admin"
|
||||
|
||||
class PermissionService:
|
||||
"""Manages role-based permissions."""
|
||||
|
||||
def __init__(self, repository: Repository):
|
||||
self.repo = repository
|
||||
|
||||
def get_permission_hierarchy(self) -> Dict[str, int]:
|
||||
return {
|
||||
Permissions.USER.value: 1,
|
||||
Permissions.GENERATOR.value: 2,
|
||||
Permissions.ADMIN.value: 3
|
||||
}
|
||||
|
||||
|
||||
|
||||
async def get_user_permission_level(self, member: Union[discord.Member, discord.User]) -> str:
|
||||
"""
|
||||
Determine the highest permission level for a user in a guild.
|
||||
Default is USER.
|
||||
Administrator permission in Discord implies ADMIN level.
|
||||
"""
|
||||
# Handle DMs or non-guild context where we have User instead of Member
|
||||
if isinstance(member, discord.User):
|
||||
return Permissions.USER.value
|
||||
|
||||
if member.guild_permissions.administrator:
|
||||
return Permissions.ADMIN.value
|
||||
|
||||
# Fetch configured roles for this server
|
||||
server_roles = await self.repo.get_server_roles(str(member.guild.id))
|
||||
|
||||
if not server_roles:
|
||||
return Permissions.USER.value
|
||||
|
||||
hierarchy = self.get_permission_hierarchy()
|
||||
current_level = Permissions.USER.value
|
||||
current_score = hierarchy[current_level]
|
||||
|
||||
# Check user's roles against configured roles
|
||||
member_role_ids = [str(r.id) for r in member.roles]
|
||||
|
||||
for server_role in server_roles:
|
||||
if server_role.role_discord_id in member_role_ids:
|
||||
level = server_role.permission_level
|
||||
if level in hierarchy and hierarchy[level] > current_score:
|
||||
current_level = level
|
||||
current_score = hierarchy[level]
|
||||
|
||||
return current_level
|
||||
|
||||
async def check_permission(self, member: discord.Member, required_level: str) -> bool:
|
||||
"""Check if user meets the required permission level."""
|
||||
user_level = await self.get_user_permission_level(member)
|
||||
hierarchy = self.get_permission_hierarchy()
|
||||
|
||||
return hierarchy.get(user_level, 0) >= hierarchy.get(required_level, 0)
|
||||
|
||||
# Helper decorator for checking permissions in commands
|
||||
def require_permission(level: str):
|
||||
async def predicate(interaction: discord.Interaction):
|
||||
if not interaction.guild:
|
||||
await interaction.response.send_message("⛔ This command cannot be used in DMs.", ephemeral=True)
|
||||
return False
|
||||
|
||||
|
||||
# We need to access the bot instance to get the permission service
|
||||
bot = interaction.client
|
||||
if not hasattr(bot, "permission_service"):
|
||||
logger.error("Bot instance missing permission_service")
|
||||
return False
|
||||
|
||||
has_perm = await bot.permission_service.check_permission(interaction.user, level)
|
||||
|
||||
if not has_perm:
|
||||
await interaction.response.send_message(
|
||||
f"⛔ You need **{level.upper()}** permission to use this command.",
|
||||
ephemeral=True
|
||||
)
|
||||
return has_perm
|
||||
return discord.app_commands.check(predicate)
|
||||
@@ -0,0 +1,48 @@
|
||||
# ComfyUI-DiscordSend Bot Configuration
|
||||
# Copy this file to bot/config.yaml and fill in your values
|
||||
# Environment variables will override these settings (see .env.example)
|
||||
|
||||
discord:
|
||||
# Your Discord bot token (required)
|
||||
# Get one at: https://discord.com/developers/applications
|
||||
token: "YOUR_BOT_TOKEN_HERE"
|
||||
|
||||
# Your Discord application ID (optional)
|
||||
# Used for slash command registration
|
||||
application_id: "YOUR_APP_ID_HERE"
|
||||
|
||||
comfyui:
|
||||
# ComfyUI server URL (default: http://127.0.0.1:8188)
|
||||
url: "http://127.0.0.1:8188"
|
||||
|
||||
# WebSocket URL for real-time updates (auto-derived from url if not set)
|
||||
# ws_url: "ws://127.0.0.1:8188/ws"
|
||||
|
||||
# Request timeout in seconds
|
||||
timeout: 30
|
||||
|
||||
defaults:
|
||||
# Maximum pending jobs per user per server
|
||||
max_queue_per_user: 3
|
||||
|
||||
# How often to update progress embeds (seconds)
|
||||
progress_update_interval: 2.0
|
||||
|
||||
# Default generation parameters
|
||||
default_steps: 20
|
||||
default_cfg: 7.0
|
||||
default_width: 512
|
||||
default_height: 512
|
||||
|
||||
# Path to default workflow (relative to bot/data/)
|
||||
# workflow_path: "default_workflow_api.json"
|
||||
|
||||
database:
|
||||
# Database URL (default: SQLite in bot/data/bot.db)
|
||||
# For PostgreSQL: postgresql+asyncpg://user:pass@localhost/dbname
|
||||
# url: "sqlite+aiosqlite:///bot/data/bot.db"
|
||||
|
||||
security:
|
||||
# Restrict bot to specific guild IDs (empty = all guilds allowed)
|
||||
allowed_guilds: []
|
||||
# Example: allowed_guilds: [123456789, 987654321]
|
||||
@@ -1,4 +0,0 @@
|
||||
{
|
||||
"count": 211,
|
||||
"uniques": 166
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:14edc1d728ca1b3dc4188432caccf1c1231ed2a16e48a5f40e4b85aa60ed2c70
|
||||
size 426901
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:3ea212a779272317e7ca77efd49e3ec3c3cda38feab41b01804e8bdd9c106e5c
|
||||
size 33919
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:e4c4024cb865e1b873d21e3b5a31af8bcb9ee937365686754d83fef38e3193f2
|
||||
size 13370
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:d60667571ab98e365bf6c8b50f79b55ca9433c7d8eee146a1c32397bc237a85f
|
||||
size 14303
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:f6801acaf97d9771969e20cb0010ba72d13e039c3f2e74a5ea189a39d1351754
|
||||
size 117428
|
||||
@@ -0,0 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:b1227f6966641c987b9943c950b09a7ba64e656ea12c852fadfd7f389ad6efa0
|
||||
size 177756
|
||||
@@ -0,0 +1,11 @@
|
||||
"""
|
||||
ComfyUI-DiscordSend Node Implementations
|
||||
|
||||
This package contains the ComfyUI custom nodes for sending media to Discord.
|
||||
"""
|
||||
|
||||
from .base_node import BaseDiscordNode
|
||||
from .image_node import DiscordSendSaveImage
|
||||
from .video_node import DiscordSendSaveVideo
|
||||
|
||||
__all__ = ['BaseDiscordNode', 'DiscordSendSaveImage', 'DiscordSendSaveVideo']
|
||||
@@ -0,0 +1,343 @@
|
||||
"""
|
||||
Base class for Discord-enabled ComfyUI nodes.
|
||||
|
||||
Provides common functionality for sending media to Discord,
|
||||
including INPUT_TYPES definitions, sanitization, and Discord integration.
|
||||
"""
|
||||
|
||||
import os
|
||||
import folder_paths
|
||||
|
||||
from shared import (
|
||||
sanitize_json_for_export,
|
||||
update_github_cdn_urls,
|
||||
send_to_discord_with_retry,
|
||||
build_filename_with_metadata,
|
||||
get_output_directory,
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
extract_cdn_urls_from_response,
|
||||
send_cdn_urls_file,
|
||||
extract_prompts_from_workflow
|
||||
)
|
||||
|
||||
|
||||
class BaseDiscordNode:
|
||||
"""
|
||||
Base class for Discord-enabled ComfyUI nodes.
|
||||
|
||||
Provides common functionality for:
|
||||
- Filename generation with metadata
|
||||
- Output directory management
|
||||
- Discord webhook integration
|
||||
- GitHub CDN URL updates
|
||||
- Workflow data sanitization
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
self.compress_level = 4
|
||||
self.output_dir = None
|
||||
|
||||
@staticmethod
|
||||
def get_discord_input_types():
|
||||
"""
|
||||
Returns Discord-related INPUT_TYPES fields.
|
||||
|
||||
These can be merged into a node's INPUT_TYPES definition.
|
||||
"""
|
||||
return {
|
||||
"send_to_discord": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If enabled, will send the media to Discord via webhook."
|
||||
}),
|
||||
"webhook_url": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "Discord webhook URL. Get this from Discord server settings > Integrations > Webhooks."
|
||||
}),
|
||||
"discord_message": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "Message to include with the media when sending to Discord."
|
||||
}),
|
||||
"include_prompts_in_message": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If enabled, will include the generation prompts in the Discord message."
|
||||
}),
|
||||
"send_workflow_json": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If enabled, will send the workflow JSON alongside the media."
|
||||
}),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_cdn_input_types():
|
||||
"""
|
||||
Returns CDN and GitHub-related INPUT_TYPES fields.
|
||||
"""
|
||||
return {
|
||||
"save_cdn_urls": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If enabled, will extract and save Discord CDN URLs."
|
||||
}),
|
||||
"github_cdn_update": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "If enabled, will update a GitHub repository with the CDN URLs."
|
||||
}),
|
||||
"github_repo": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "GitHub repository in format 'username/repo'."
|
||||
}),
|
||||
"github_token": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": False,
|
||||
"tooltip": "GitHub Personal Access Token (PAT). Generate at: Settings > Developer settings > Tokens. \n⚠️ Requires 'repo' scope (or 'public_repo') to upload files."
|
||||
}),
|
||||
"github_file_path": ("STRING", {
|
||||
"default": "cdn_urls.md",
|
||||
"multiline": False,
|
||||
"tooltip": "Path to the file in the GitHub repository to update."
|
||||
}),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_filename_input_types(
|
||||
add_date_default: bool = False,
|
||||
add_time_default: bool = True,
|
||||
add_dimensions_default: bool = False
|
||||
):
|
||||
"""
|
||||
Returns filename metadata INPUT_TYPES fields.
|
||||
|
||||
Args:
|
||||
add_date_default: Default value for add_date
|
||||
add_time_default: Default value for add_time
|
||||
add_dimensions_default: Default value for add_dimensions
|
||||
"""
|
||||
return {
|
||||
"add_date": ("BOOLEAN", {
|
||||
"default": add_date_default,
|
||||
"tooltip": "Add date (YYYY-MM-DD) to the filename."
|
||||
}),
|
||||
"add_time": ("BOOLEAN", {
|
||||
"default": add_time_default,
|
||||
"tooltip": "Add time (HH-MM-SS) to the filename."
|
||||
}),
|
||||
"add_dimensions": ("BOOLEAN", {
|
||||
"default": add_dimensions_default,
|
||||
"tooltip": "Add dimensions (WxH) to the filename."
|
||||
}),
|
||||
}
|
||||
|
||||
def sanitize_workflow_data(self, prompt, extra_pnginfo):
|
||||
"""
|
||||
Sanitize workflow data by removing sensitive information.
|
||||
|
||||
Args:
|
||||
prompt: The prompt data
|
||||
extra_pnginfo: Extra PNG info including workflow
|
||||
|
||||
Returns:
|
||||
Tuple of (sanitized_prompt, sanitized_extra_pnginfo,
|
||||
original_prompt, original_extra_pnginfo)
|
||||
"""
|
||||
# Store original references for prompt extraction
|
||||
original_prompt = prompt
|
||||
original_extra_pnginfo = extra_pnginfo
|
||||
|
||||
# Sanitize workflow data
|
||||
if prompt is not None:
|
||||
prompt = sanitize_json_for_export(prompt)
|
||||
|
||||
if extra_pnginfo is not None:
|
||||
extra_pnginfo = sanitize_json_for_export(extra_pnginfo)
|
||||
|
||||
return prompt, extra_pnginfo, original_prompt, original_extra_pnginfo
|
||||
|
||||
def build_filename_prefix(
|
||||
self,
|
||||
filename_prefix: str,
|
||||
add_date: bool,
|
||||
add_time: bool,
|
||||
add_dimensions: bool = False,
|
||||
width: int = None,
|
||||
height: int = None
|
||||
):
|
||||
"""
|
||||
Build filename prefix with metadata.
|
||||
|
||||
Returns:
|
||||
Tuple of (modified_prefix, info_dict)
|
||||
"""
|
||||
info_dict = {}
|
||||
filename_prefix, info_dict = build_filename_with_metadata(
|
||||
prefix=filename_prefix,
|
||||
add_date=add_date,
|
||||
add_time=add_time,
|
||||
add_dimensions=add_dimensions,
|
||||
width=width,
|
||||
height=height,
|
||||
info_dict=info_dict
|
||||
)
|
||||
filename_prefix += self.prefix_append
|
||||
return filename_prefix, info_dict
|
||||
|
||||
def get_dest_folder(self, save_output: bool):
|
||||
"""
|
||||
Get the destination folder for output files.
|
||||
|
||||
Args:
|
||||
save_output: Whether to save to output directory (True) or temp (False)
|
||||
|
||||
Returns:
|
||||
Path to the destination folder
|
||||
"""
|
||||
return get_output_directory(
|
||||
save_output=save_output,
|
||||
comfy_output_dir=folder_paths.get_output_directory(),
|
||||
temp_dir=folder_paths.get_temp_directory()
|
||||
)
|
||||
|
||||
def extract_workflow_from_metadata(self, original_prompt, original_extra_pnginfo):
|
||||
"""
|
||||
Extract workflow data from metadata.
|
||||
|
||||
Args:
|
||||
original_prompt: Original prompt data
|
||||
original_extra_pnginfo: Original extra PNG info
|
||||
|
||||
Returns:
|
||||
Workflow data dict, or None if not found
|
||||
"""
|
||||
workflow_data = None
|
||||
|
||||
# First try to get workflow from extra_pnginfo
|
||||
if (original_extra_pnginfo is not None and
|
||||
isinstance(original_extra_pnginfo, dict) and
|
||||
"workflow" in original_extra_pnginfo):
|
||||
workflow_data = original_extra_pnginfo["workflow"]
|
||||
|
||||
# If no workflow in extra_pnginfo, check if prompt is actually a workflow
|
||||
if workflow_data is None and original_prompt is not None:
|
||||
if isinstance(original_prompt, dict) and "nodes" in original_prompt:
|
||||
workflow_data = original_prompt
|
||||
|
||||
return workflow_data
|
||||
|
||||
def build_prompt_message(self, workflow_data):
|
||||
"""
|
||||
Extract and build prompt message from workflow data.
|
||||
|
||||
Args:
|
||||
workflow_data: Workflow data dict
|
||||
|
||||
Returns:
|
||||
Formatted prompt section string, or empty string
|
||||
"""
|
||||
if workflow_data is None:
|
||||
return ""
|
||||
|
||||
positive_prompt, negative_prompt = extract_prompts_from_workflow(workflow_data)
|
||||
return build_prompt_section(positive_prompt, negative_prompt)
|
||||
|
||||
def send_discord_files(
|
||||
self,
|
||||
webhook_url: str,
|
||||
files: dict,
|
||||
data: dict,
|
||||
save_cdn_urls: bool = False
|
||||
):
|
||||
"""
|
||||
Send files to Discord via webhook.
|
||||
|
||||
Args:
|
||||
webhook_url: Discord webhook URL
|
||||
files: Files dict for the request
|
||||
data: Data dict for the request
|
||||
save_cdn_urls: Whether to extract CDN URLs from response
|
||||
|
||||
Returns:
|
||||
Tuple of (success, response, cdn_urls)
|
||||
"""
|
||||
cdn_urls = []
|
||||
|
||||
response = send_to_discord_with_retry(
|
||||
webhook_url,
|
||||
files=files,
|
||||
data=data
|
||||
)
|
||||
|
||||
success = response.status_code in [200, 204]
|
||||
|
||||
if success and save_cdn_urls:
|
||||
cdn_urls = extract_cdn_urls_from_response(response)
|
||||
|
||||
return success, response, cdn_urls
|
||||
|
||||
def send_cdn_urls_to_discord(
|
||||
self,
|
||||
webhook_url: str,
|
||||
cdn_urls: list,
|
||||
message: str = "Discord CDN URLs:"
|
||||
):
|
||||
"""
|
||||
Send CDN URLs as a text file to Discord.
|
||||
|
||||
Args:
|
||||
webhook_url: Discord webhook URL
|
||||
cdn_urls: List of (filename, url) tuples
|
||||
message: Message to accompany the file
|
||||
|
||||
Returns:
|
||||
True if successful
|
||||
"""
|
||||
if not cdn_urls:
|
||||
return False
|
||||
|
||||
return send_cdn_urls_file(
|
||||
webhook_url=webhook_url,
|
||||
urls=cdn_urls,
|
||||
send_func=send_to_discord_with_retry,
|
||||
message=message
|
||||
)
|
||||
|
||||
def update_github_cdn(
|
||||
self,
|
||||
cdn_urls: list,
|
||||
github_repo: str,
|
||||
github_token: str,
|
||||
github_file_path: str
|
||||
):
|
||||
"""
|
||||
Update GitHub repository with CDN URLs.
|
||||
|
||||
Args:
|
||||
cdn_urls: List of (filename, url) tuples
|
||||
github_repo: GitHub repository in 'owner/repo' format
|
||||
github_token: GitHub personal access token
|
||||
github_file_path: Path to file in repository
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message)
|
||||
"""
|
||||
if not cdn_urls:
|
||||
return False, "No CDN URLs to update"
|
||||
|
||||
print(f"Updating GitHub repository {github_repo} with {len(cdn_urls)} CDN URLs...")
|
||||
|
||||
success, message = update_github_cdn_urls(
|
||||
github_repo=github_repo,
|
||||
github_token=github_token,
|
||||
file_path=github_file_path,
|
||||
cdn_urls=cdn_urls
|
||||
)
|
||||
|
||||
if success:
|
||||
print(f"GitHub update successful: {message}")
|
||||
else:
|
||||
print(f"GitHub update failed: {message}")
|
||||
|
||||
return success, message
|
||||
@@ -0,0 +1,560 @@
|
||||
"""ComfyUI node for sending images to Discord and saving them locally."""
|
||||
|
||||
import os
|
||||
import json
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import folder_paths
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
from comfy.cli_args import args
|
||||
import re
|
||||
import cv2
|
||||
from io import BytesIO
|
||||
from uuid import uuid4
|
||||
from typing import Any, Union, List, Optional
|
||||
|
||||
# Import shared utilities
|
||||
from shared import (
|
||||
sanitize_token_from_text,
|
||||
process_batched_images,
|
||||
build_metadata_section,
|
||||
validate_path_is_safe,
|
||||
sanitize_json_for_export
|
||||
)
|
||||
|
||||
|
||||
from .base_node import BaseDiscordNode
|
||||
|
||||
|
||||
class DiscordSendSaveImage(BaseDiscordNode):
|
||||
"""
|
||||
A ComfyUI node that can send images to Discord and save them with advanced options.
|
||||
Images can be sent to Discord via webhook integration, while providing flexible
|
||||
saving options with customizable naming conventions and format options.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
self.compress_level = 4
|
||||
self.output_dir = None
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
# Get base inputs from BaseDiscordNode
|
||||
base_inputs = BaseDiscordNode.get_discord_input_types()
|
||||
cdn_inputs = BaseDiscordNode.get_cdn_input_types()
|
||||
filename_inputs = BaseDiscordNode.get_filename_input_types(add_date_default=False)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": "The images to save and/or send to Discord."}),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI-Image", "tooltip": "The prefix for the saved files. Supports %batch_num% placeholder for batch indexing."}),
|
||||
"overwrite_last": ("BOOLEAN", {"default": False, "tooltip": "⚠️ CAUTION: If enabled, new saves will REPLACE the previous file with the same name. Useful for iterative testing, dangerous for batch production. Note: You must also disable 'add_time' and 'add_date' to ensure filenames are identical."})
|
||||
},
|
||||
"optional": {
|
||||
"file_format": (["png", "jpeg", "webp"], {
|
||||
"default": "png",
|
||||
"tooltip": "The format to save images in. PNG is lossless but larger. JPEG and WebP are smaller but lossy."
|
||||
}),
|
||||
"quality": ("INT", {
|
||||
"default": 95,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"step": 1,
|
||||
"tooltip": "Quality (1-100) for JPEG/WebP. Ignored for PNG. Higher values = better quality but larger file size."
|
||||
}),
|
||||
"lossless": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Use lossless compression for WebP (PNG is always lossless). For JPEG, forces maximum quality (100)."
|
||||
}),
|
||||
"save_output": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Whether to save images to disk. When disabled, images will only be previewed in the UI."
|
||||
}),
|
||||
"show_preview": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Whether to show image previews in the UI. Disable to reduce UI clutter for large batches."
|
||||
}),
|
||||
"resize_to_power_of_2": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Resize images to nearest power of 2 dimensions (useful for game textures). ⚠️ May distort aspect ratio. Uses the algorithm selected in 'resize_method'."
|
||||
}),
|
||||
"resize_method": (["nearest-exact", "bilinear", "bicubic", "lanczos", "box"], {
|
||||
"default": "lanczos",
|
||||
"tooltip": "Resampling algorithm used ONLY when 'resize_to_power_of_2' is enabled. Ignored otherwise. \n• lanczos: Best for photos\n• nearest-exact: Best for pixel art\n• bilinear/bicubic: Faster"
|
||||
}),
|
||||
"include_format_in_message": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Whether to include the image format in the Discord message."
|
||||
}),
|
||||
"group_batched_images": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Group all images from a batch into a single Discord message with a gallery, rather than sending each one separately. Maximum is 9 images."
|
||||
}),
|
||||
# Mix in shared options
|
||||
**filename_inputs,
|
||||
**base_inputs,
|
||||
**cdn_inputs,
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("image_path",)
|
||||
FUNCTION = "save_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "image/output"
|
||||
DESCRIPTION = "Saves images with advanced options and can send them to Discord via webhook integration. Returns the path to the first saved image."
|
||||
|
||||
@classmethod
|
||||
def CONTEXT_MENUS(s):
|
||||
return {
|
||||
"Show Preview": lambda self, **kwargs: {"show_preview": True},
|
||||
"Hide Preview": lambda self, **kwargs: {"show_preview": False},
|
||||
}
|
||||
|
||||
def save_images(self, images, filename_prefix="ComfyUI-Image", overwrite_last=False,
|
||||
file_format="png", quality=95, lossless=True, add_date=False, add_time=False,
|
||||
add_dimensions=False, resize_to_power_of_2=False, save_output=True,
|
||||
resize_method="lanczos", show_preview=True, send_to_discord=False, webhook_url="", discord_message="",
|
||||
include_prompts_in_message=False, include_format_in_message=False, send_workflow_json=False,
|
||||
group_batched_images=True, save_cdn_urls=False, github_cdn_update=False, github_repo="",
|
||||
github_token="", github_file_path="cdn_urls.md", prompt=None, extra_pnginfo=None):
|
||||
"""
|
||||
Save images for and optionally send to Discord.
|
||||
"""
|
||||
results = []
|
||||
output_files = []
|
||||
discord_sent_files = []
|
||||
discord_send_success = True
|
||||
|
||||
# For batch grouping
|
||||
batch_discord_files = []
|
||||
batch_discord_data = {}
|
||||
batch_workflow_json = None
|
||||
|
||||
# For tracking Discord CDN URLs
|
||||
discord_cdn_urls = []
|
||||
batch_cdn_urls = []
|
||||
|
||||
# 1. Sanitize workflow data using base class method
|
||||
prompt, extra_pnginfo, original_prompt, original_extra_pnginfo = self.sanitize_workflow_data(
|
||||
prompt, extra_pnginfo
|
||||
)
|
||||
|
||||
# 2. Build filename prefix with metadata using base class method
|
||||
filename_prefix, image_info = self.build_filename_prefix(
|
||||
filename_prefix, add_date, add_time, False, None, None
|
||||
)
|
||||
|
||||
# Add prefix append
|
||||
filename_prefix += self.prefix_append
|
||||
|
||||
# 3. Get output directory using base class method
|
||||
dest_folder = self.get_dest_folder(save_output)
|
||||
|
||||
# Setup paths using ComfyUI's path validation
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix, dest_folder, images[0].shape[1], images[0].shape[0])
|
||||
|
||||
# For overwrite functionality, we'll just always use the same counter instead of bypassing validation
|
||||
if overwrite_last:
|
||||
counter = 1 # Always use the same counter value for overwriting
|
||||
else:
|
||||
# When not overwriting, we need to find the highest existing counter and start from there
|
||||
# This ensures we're always creating new files
|
||||
try:
|
||||
# Get all existing files with this prefix
|
||||
base_filename = os.path.basename(filename).replace("%batch_num%", "")
|
||||
existing_files = [f for f in os.listdir(full_output_folder)
|
||||
if os.path.basename(f).startswith(base_filename)]
|
||||
|
||||
if existing_files:
|
||||
# Extract counters from filenames
|
||||
existing_counters = []
|
||||
for f in existing_files:
|
||||
# Extract counter pattern (5 digits) from filename
|
||||
counter_match = re.search(r'_(\d{5})\.', f)
|
||||
if counter_match:
|
||||
existing_counters.append(int(counter_match.group(1)))
|
||||
|
||||
# Also try alternative pattern where the counter is followed by extension
|
||||
counter_match = re.search(r'_(\d{5})_\.', f)
|
||||
if counter_match:
|
||||
existing_counters.append(int(counter_match.group(1)))
|
||||
|
||||
# Set counter to one more than the highest existing counter
|
||||
if existing_counters:
|
||||
counter = max(existing_counters) + 1
|
||||
except Exception as e:
|
||||
print(f"Error determining next file counter: {e}")
|
||||
# Default to ComfyUI's counter if we can't determine the next one
|
||||
|
||||
print(f"Using counter: {counter} for {'overwriting' if overwrite_last else 'new files'}")
|
||||
print(f"Output prefix: {filename_prefix}")
|
||||
|
||||
# Map resize method strings to PIL resize methods
|
||||
resize_methods = {
|
||||
"nearest-exact": Image.NEAREST,
|
||||
"bilinear": Image.BILINEAR,
|
||||
"bicubic": Image.BICUBIC,
|
||||
"lanczos": Image.LANCZOS,
|
||||
"box": Image.BOX
|
||||
}
|
||||
|
||||
# Handle different versions of PIL
|
||||
if hasattr(Image, 'Resampling'):
|
||||
resize_methods = {
|
||||
"nearest-exact": Image.Resampling.NEAREST,
|
||||
"bilinear": Image.Resampling.BILINEAR,
|
||||
"bicubic": Image.Resampling.BICUBIC,
|
||||
"lanczos": Image.Resampling.LANCZOS,
|
||||
"box": Image.Resampling.BOX
|
||||
}
|
||||
|
||||
# Get the selected resize method, default to LANCZOS if not found
|
||||
selected_resize_method = resize_methods.get(resize_method, Image.LANCZOS)
|
||||
|
||||
# Initialize Discord sender if enabled
|
||||
discord_success = False
|
||||
if send_to_discord and webhook_url:
|
||||
print(f"Discord integration enabled, preparing to send images to webhook")
|
||||
discord_success = True # Will be set to False if any send fails
|
||||
|
||||
# Initialize message_prefix for all Discord messages
|
||||
# This ensures prompts have a place to be attached regardless of other options
|
||||
image_info["message_prefix"] = ""
|
||||
|
||||
elif send_to_discord and not webhook_url:
|
||||
print("Discord integration was enabled but no webhook URL was provided")
|
||||
|
||||
# Build image info message using shared utility
|
||||
if send_to_discord and webhook_url and (add_date or add_time or add_dimensions or resize_to_power_of_2 or include_format_in_message):
|
||||
info_message = build_metadata_section(
|
||||
info_dict=image_info,
|
||||
include_date=add_date,
|
||||
include_time=add_time,
|
||||
include_dimensions=False, # Dimensions added later after processing
|
||||
include_format=include_format_in_message,
|
||||
file_format=file_format,
|
||||
section_title="Image Information"
|
||||
)
|
||||
image_info["message_prefix"] = info_message
|
||||
print("Prepared image information section for Discord message")
|
||||
|
||||
# 4. Extract and build prompts section
|
||||
if send_to_discord and include_prompts_in_message:
|
||||
workflow_data = self.extract_workflow_from_metadata(original_prompt, original_extra_pnginfo)
|
||||
if workflow_data:
|
||||
prompt_message = self.build_prompt_message(workflow_data)
|
||||
if prompt_message:
|
||||
image_info["prompt_message"] = prompt_message
|
||||
print("Prepared prompts for Discord message")
|
||||
|
||||
# Optimization: Create metadata once for the entire batch
|
||||
# This prevents redundant sanitization and JSON serialization for every image
|
||||
metadata = None
|
||||
if not args.disable_metadata:
|
||||
metadata = PngInfo()
|
||||
if prompt is not None:
|
||||
# Prompt is already sanitized at start of function
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
# extra_pnginfo is already sanitized at start of function
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
batch_counter = 0
|
||||
for chunk in process_batched_images(images):
|
||||
if len(chunk.shape) == 4:
|
||||
chunk_images = [chunk[i] for i in range(chunk.shape[0])]
|
||||
else:
|
||||
chunk_images = [chunk]
|
||||
|
||||
for image_np in chunk_images:
|
||||
batch_number = batch_counter
|
||||
batch_counter += 1
|
||||
# Convert the tensor to a PIL image
|
||||
i = image_np
|
||||
img = Image.fromarray(i)
|
||||
|
||||
# Track if resizing happened to optimize Discord encoding later
|
||||
was_resized = False
|
||||
orig_width, orig_height = img.size
|
||||
|
||||
# Resize to power of 2 if enabled
|
||||
if resize_to_power_of_2:
|
||||
new_width = 2 ** int(np.log2(orig_width) + 0.5)
|
||||
new_height = 2 ** int(np.log2(orig_height) + 0.5)
|
||||
|
||||
print(f"Resizing image from {orig_width}x{orig_height} to {new_width}x{new_height} (power of 2)")
|
||||
|
||||
if send_to_discord and webhook_url and batch_number == 0:
|
||||
image_info["original_dimensions"] = f"{orig_width}x{orig_height}"
|
||||
image_info["resized_dimensions"] = f"{new_width}x{new_height}"
|
||||
|
||||
if (new_width != orig_width or new_height != orig_height):
|
||||
try:
|
||||
img = img.resize((new_width, new_height), selected_resize_method)
|
||||
was_resized = True
|
||||
print(f"Successfully resized using {resize_method} method")
|
||||
except Exception as e:
|
||||
print(f"Error during power of 2 resize: {e}")
|
||||
img = img.resize((new_width, new_height), Image.BICUBIC)
|
||||
was_resized = True
|
||||
print("Fallback to BICUBIC resize method due to error")
|
||||
|
||||
# Get dimensions
|
||||
width, height = img.size
|
||||
|
||||
# Add dimensions to filename if enabled
|
||||
dimensions_suffix = ""
|
||||
if add_dimensions:
|
||||
dimensions_suffix = f"_{width}x{height}"
|
||||
filename_prefix += dimensions_suffix
|
||||
|
||||
if send_to_discord and webhook_url and batch_number == 0:
|
||||
image_info["dimensions"] = f"{width}x{height}"
|
||||
|
||||
# Add image information to Discord message if this is the first image
|
||||
if send_to_discord and webhook_url and batch_number == 0:
|
||||
if "message_prefix" in image_info:
|
||||
info_message = image_info["message_prefix"]
|
||||
has_resize_dimensions = "original_dimensions" in image_info and "resized_dimensions" in image_info
|
||||
has_dimensions = "dimensions" in image_info
|
||||
|
||||
if (has_resize_dimensions or has_dimensions) and not info_message:
|
||||
info_message = "\n\n**Image Information:**\n"
|
||||
|
||||
if has_resize_dimensions:
|
||||
info_message += f"**Original Dimensions:** {image_info['original_dimensions']}\n"
|
||||
info_message += f"**Resized Dimensions:** {image_info['resized_dimensions']} (Power of 2)\n"
|
||||
elif has_dimensions:
|
||||
info_message += f"**Dimensions:** {image_info['dimensions']}\n"
|
||||
|
||||
if info_message:
|
||||
discord_message += info_message
|
||||
print("Added image information to Discord message")
|
||||
|
||||
if "prompt_message" in image_info:
|
||||
discord_message += image_info["prompt_message"]
|
||||
print("Added prompts to Discord message after image information")
|
||||
|
||||
# For Discord output
|
||||
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
||||
|
||||
# Add dimensions tag before the counter if enabled
|
||||
if add_dimensions and dimensions_suffix not in filename_with_batch_num:
|
||||
base_name = os.path.splitext(filename_with_batch_num)[0]
|
||||
filename_with_batch_num = f"{base_name}{dimensions_suffix}"
|
||||
|
||||
# File extension based on format
|
||||
extension = f".{file_format}"
|
||||
file = f"{filename_with_batch_num}_{counter:05}_{extension}"
|
||||
|
||||
if file.endswith(f"_{extension}"):
|
||||
file = file[:-len(f"_{extension}")] + extension
|
||||
|
||||
filepath = os.path.join(full_output_folder, file)
|
||||
|
||||
# Security: Validate output path to prevent symlink overwrites
|
||||
validate_path_is_safe(filepath, base_dir=full_output_folder)
|
||||
|
||||
try:
|
||||
# Save the image based on format
|
||||
if file_format == "png":
|
||||
img.save(filepath, pnginfo=metadata, compress_level=self.compress_level)
|
||||
elif file_format == "jpeg":
|
||||
jpeg_quality = 100 if lossless else quality
|
||||
img.save(filepath, format="JPEG", quality=jpeg_quality)
|
||||
elif file_format == "webp":
|
||||
if lossless:
|
||||
img.save(filepath, format="WEBP", lossless=True)
|
||||
else:
|
||||
img.save(filepath, format="WEBP", quality=quality)
|
||||
|
||||
output_files.append(filepath)
|
||||
|
||||
print(f"Saved image with dimensions: {img.size[0]}x{img.size[1]}")
|
||||
|
||||
results.append({
|
||||
"filename": file,
|
||||
"subfolder": "discord_output/" + (subfolder if subfolder else "") if save_output else "",
|
||||
"type": "output" if save_output else "temp",
|
||||
"path": filepath
|
||||
})
|
||||
|
||||
# Send to Discord if enabled
|
||||
if send_to_discord and webhook_url:
|
||||
try:
|
||||
discord_filename = f"{uuid4()}.{file_format}"
|
||||
file_bytes = BytesIO()
|
||||
|
||||
if file_format == "jpeg":
|
||||
save_img = img
|
||||
if save_img.mode == 'RGBA':
|
||||
save_img = save_img.convert('RGB')
|
||||
jpeg_quality = 100 if lossless else quality
|
||||
save_img.save(file_bytes, format="JPEG", quality=jpeg_quality)
|
||||
file_bytes.seek(0)
|
||||
|
||||
elif file_format == "png":
|
||||
if not was_resized:
|
||||
img_cv = i
|
||||
else:
|
||||
img_cv = np.array(img)
|
||||
|
||||
if len(img_cv.shape) == 3 and img_cv.shape[2] == 3:
|
||||
img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
|
||||
if len(img_cv.shape) == 2:
|
||||
img_cv = cv2.cvtColor(img_cv, cv2.COLOR_GRAY2BGR)
|
||||
elif len(img_cv.shape) == 3 and img_cv.shape[2] == 4:
|
||||
img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGBA2BGRA)
|
||||
|
||||
_, buffer = cv2.imencode('.png', img_cv)
|
||||
file_bytes = BytesIO(buffer)
|
||||
|
||||
elif file_format == "webp":
|
||||
try:
|
||||
if lossless:
|
||||
img.save(file_bytes, format="WEBP", lossless=True)
|
||||
else:
|
||||
img.save(file_bytes, format="WEBP", quality=quality)
|
||||
file_bytes.seek(0)
|
||||
except Exception as e:
|
||||
print(f"Error with WebP encoding for Discord: {e}, falling back to PNG")
|
||||
discord_filename = f"{os.path.splitext(discord_filename)[0]}.png"
|
||||
file_bytes = BytesIO() # Reset buffer
|
||||
img.save(file_bytes, format="PNG", compress_level=self.compress_level)
|
||||
file_bytes.seek(0)
|
||||
|
||||
if group_batched_images:
|
||||
batch_discord_files.append((discord_filename, file_bytes.getvalue()))
|
||||
|
||||
# Prepare workflow JSON only once for the whole batch
|
||||
if batch_number == 0 and send_workflow_json and (prompt is not None or extra_pnginfo is not None):
|
||||
wflow = self.extract_workflow_from_metadata(original_prompt, original_extra_pnginfo)
|
||||
if wflow:
|
||||
batch_workflow_json = wflow
|
||||
|
||||
if batch_number == 0 and discord_message:
|
||||
batch_discord_data["content"] = discord_message
|
||||
|
||||
else:
|
||||
# Immediate send
|
||||
files = {
|
||||
"file": (discord_filename, file_bytes.getvalue())
|
||||
}
|
||||
|
||||
if send_workflow_json and (prompt is not None or extra_pnginfo is not None):
|
||||
wflow = self.extract_workflow_from_metadata(original_prompt, original_extra_pnginfo)
|
||||
if wflow:
|
||||
# Sanitize to remove webhook URLs and tokens
|
||||
wflow = sanitize_json_for_export(wflow)
|
||||
json_filename = f"{os.path.splitext(discord_filename)[0]}.json"
|
||||
files["workflow"] = (json_filename, json.dumps(wflow, indent=2).encode('utf-8'))
|
||||
|
||||
data = {}
|
||||
if discord_message:
|
||||
data["content"] = discord_message
|
||||
|
||||
success, response, new_urls = self.send_discord_files(webhook_url, files, data, save_cdn_urls)
|
||||
|
||||
if success:
|
||||
print(f"Successfully sent image {batch_number+1} to Discord")
|
||||
discord_sent_files.append(discord_filename)
|
||||
if new_urls:
|
||||
batch_cdn_urls.extend(new_urls)
|
||||
self.send_cdn_urls_to_discord(webhook_url, new_urls, "Discord CDN URLs for the uploaded images:")
|
||||
else:
|
||||
print(f"Error: Discord returned status code {response.status_code}")
|
||||
discord_send_success = False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing image for Discord: {e}")
|
||||
discord_send_success = False
|
||||
|
||||
if not overwrite_last:
|
||||
counter += 1
|
||||
except Exception as e:
|
||||
print(f"Error saving image: {e}")
|
||||
|
||||
if results:
|
||||
if save_output:
|
||||
print(f"DiscordSendSaveImage: Saved {len(results)} images to {full_output_folder}")
|
||||
else:
|
||||
print("DiscordSendSaveImage: Preview only mode - no images saved to disk")
|
||||
|
||||
if send_to_discord and discord_sent_files:
|
||||
print("DiscordSendSaveImage: Successfully sent all images to Discord")
|
||||
elif send_to_discord and not discord_send_success:
|
||||
print("DiscordSendSaveImage: There were errors sending some images to Discord")
|
||||
else:
|
||||
print("DiscordSendSaveImage: No images were processed")
|
||||
|
||||
# Send batch to Discord
|
||||
if send_to_discord and webhook_url and group_batched_images and batch_discord_files:
|
||||
try:
|
||||
print(f"Sending {len(batch_discord_files)} images as a batch to Discord...")
|
||||
|
||||
files = {}
|
||||
for i, (filename, file_bytes) in enumerate(batch_discord_files):
|
||||
files[f"file{i}"] = (filename, file_bytes)
|
||||
|
||||
if send_workflow_json and batch_workflow_json:
|
||||
# Sanitize to remove webhook URLs and tokens
|
||||
sanitized_workflow = sanitize_json_for_export(batch_workflow_json)
|
||||
json_filename = f"workflow-{uuid4()}.json"
|
||||
json_data = json.dumps(sanitized_workflow, indent=2)
|
||||
files["workflow"] = (json_filename, json_data.encode('utf-8'))
|
||||
|
||||
success, response, new_urls = self.send_discord_files(webhook_url, files, batch_discord_data, save_cdn_urls)
|
||||
|
||||
if success:
|
||||
print(f"Successfully sent batch of {len(batch_discord_files)} images to Discord as a gallery")
|
||||
discord_sent_files = ["batch_gallery"]
|
||||
if save_cdn_urls and new_urls:
|
||||
batch_cdn_urls.extend(new_urls)
|
||||
self.send_cdn_urls_to_discord(webhook_url, new_urls, "Discord CDN URLs for the uploaded images:")
|
||||
else:
|
||||
error_msg = sanitize_token_from_text(response.text, webhook_url)
|
||||
print(f"Error sending batch to Discord: Status code {response.status_code} - {error_msg}")
|
||||
discord_send_success = False
|
||||
except Exception as e:
|
||||
print(f"Error sending batch to Discord: {e}")
|
||||
discord_send_success = False
|
||||
|
||||
# Update GitHub repository
|
||||
if github_cdn_update and send_to_discord and (discord_cdn_urls or batch_cdn_urls):
|
||||
urls_to_send = discord_cdn_urls if discord_cdn_urls else batch_cdn_urls
|
||||
self.update_github_cdn(urls_to_send, github_repo, github_token, github_file_path)
|
||||
|
||||
elif github_cdn_update:
|
||||
reasons = []
|
||||
if not send_to_discord: reasons.append("send_to_discord is disabled")
|
||||
if not (discord_cdn_urls or batch_cdn_urls): reasons.append("no CDN URLs were collected")
|
||||
if not github_repo: reasons.append("github_repo is empty")
|
||||
if not github_token: reasons.append("github_token is empty")
|
||||
if not github_file_path: reasons.append("github_file_path is empty")
|
||||
print(f"GitHub update was enabled but not triggered because: {', '.join(reasons)}")
|
||||
|
||||
# Return results
|
||||
if show_preview:
|
||||
return {"ui": {"images": results}, "result": ((save_output, output_files, discord_send_success if send_to_discord else None),)}, output_files[0] if output_files else ""
|
||||
else:
|
||||
return {"ui": {}, "result": ((save_output, output_files, discord_send_success if send_to_discord else None),)}, output_files[0] if output_files else ""
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, images, filename_prefix="ComfyUI-Image", overwrite_last=False,
|
||||
file_format="png", quality=95, lossless=True, add_date=False, add_time=False,
|
||||
add_dimensions=False, resize_to_power_of_2=False, save_output=True,
|
||||
resize_method="lanczos", show_preview=True, send_to_discord=False, webhook_url="", discord_message="",
|
||||
include_prompts_in_message=False, include_format_in_message=False, group_batched_images=True,
|
||||
send_workflow_json=False, save_cdn_urls=False, github_cdn_update=False, github_repo="",
|
||||
github_token="", github_file_path="cdn_urls.md", prompt=None, extra_pnginfo=None):
|
||||
return True
|
||||
+1070
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,15 @@
|
||||
[project]
|
||||
name = "comfyui-discordsend"
|
||||
description = "A ComfyUI extension that enables seamless sharing of AI-generated images and videos directly to Discord."
|
||||
version = "2.0.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["requests>=2.25.0"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/AEmotionStudio/ComfyUI-DiscordSend"
|
||||
# Used by Comfy Registry https://comfyregistry.org
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "aemotionstudio"
|
||||
DisplayName = "ComfyUI-DiscordSend"
|
||||
Icon = ""
|
||||
@@ -0,0 +1,24 @@
|
||||
# Discord Bot Dependencies
|
||||
# Install with: pip install -r requirements-bot.txt
|
||||
#
|
||||
# Note: This includes node dependencies plus bot-specific packages
|
||||
|
||||
# Shared with nodes
|
||||
requests>=2.25.0
|
||||
Pillow>=9.0.0
|
||||
numpy>=1.20.0
|
||||
|
||||
# Discord bot framework
|
||||
discord.py>=2.3.0
|
||||
|
||||
# Async HTTP client for ComfyUI API
|
||||
aiohttp>=3.9.0
|
||||
|
||||
# Database ORM for job tracking
|
||||
sqlalchemy>=2.0.0
|
||||
|
||||
# Async SQLite driver
|
||||
aiosqlite>=0.19.0
|
||||
|
||||
# YAML config parsing
|
||||
pyyaml>=6.0.0
|
||||
@@ -0,0 +1,6 @@
|
||||
# Node-only Dependencies (same as requirements.txt)
|
||||
# This file exists for clarity - use requirements.txt for ComfyUI Manager
|
||||
|
||||
requests>=2.25.0
|
||||
Pillow>=9.0.0
|
||||
numpy>=1.20.0
|
||||
@@ -0,0 +1,10 @@
|
||||
# ComfyUI-DiscordSend Node Dependencies
|
||||
#
|
||||
# This file contains minimal dependencies for the ComfyUI nodes only.
|
||||
# ComfyUI Manager will automatically install these.
|
||||
#
|
||||
# For Discord bot support, manually run: pip install -r requirements-bot.txt
|
||||
|
||||
requests>=2.25.0
|
||||
Pillow>=9.0.0
|
||||
numpy>=1.20.0
|
||||
Executable
+42
@@ -0,0 +1,42 @@
|
||||
#!/bin/bash
|
||||
|
||||
# auto_review.sh - Automatically review the latest changes using Gemini CLI
|
||||
|
||||
# 1. Get the latest changes (staged or last commit)
|
||||
# If there are staged changes, review them. Otherwise review last commit.
|
||||
if git diff --quiet --cached; then
|
||||
# No staged changes, check last commit
|
||||
DIFF_CONTENT=$(git show HEAD)
|
||||
CONTEXT="Review the following changes from the last commit:"
|
||||
else
|
||||
# Staged changes exist
|
||||
DIFF_CONTENT=$(git diff --cached)
|
||||
CONTEXT="Review the following staged changes:"
|
||||
fi
|
||||
|
||||
if [ -z "$DIFF_CONTENT" ]; then
|
||||
echo "No changes found to review."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# 2. Construct Prompt
|
||||
PROMPT="You are a Senior Software Engineer acting as a code reviewer.
|
||||
Review the following code changes for:
|
||||
1. Potential bugs or race conditions
|
||||
2. Security vulnerabilities
|
||||
3. Code style and best practices
|
||||
4. Logical errors
|
||||
|
||||
Be concise and constructive.
|
||||
|
||||
$CONTEXT
|
||||
\`\`\`diff
|
||||
$DIFF_CONTENT
|
||||
\`\`\`
|
||||
"
|
||||
|
||||
# 3. Call Gemini CLI
|
||||
echo "🤖 Asking Gemini to review changes..."
|
||||
echo "----------------------------------------"
|
||||
gemini "$PROMPT"
|
||||
echo "----------------------------------------"
|
||||
@@ -0,0 +1,70 @@
|
||||
"""
|
||||
ComfyUI-DiscordSend Shared Utilities
|
||||
|
||||
This package contains shared utilities used by both ComfyUI nodes and the Discord bot.
|
||||
Organized into subpackages:
|
||||
- discord: Discord webhook and message utilities
|
||||
- media: Image and video processing utilities
|
||||
- workflow: ComfyUI workflow manipulation utilities
|
||||
"""
|
||||
|
||||
# Re-export commonly used utilities for convenience
|
||||
from .workflow.sanitizer import sanitize_json_for_export
|
||||
from .workflow.prompt_extractor import extract_prompts_from_workflow
|
||||
from .workflow.workflow_builder import WorkflowBuilder
|
||||
from .discord.webhook_client import (
|
||||
DiscordWebhookClient,
|
||||
validate_webhook_url,
|
||||
send_to_discord_with_retry,
|
||||
sanitize_token_from_text
|
||||
)
|
||||
from .discord.message_builder import (
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
build_discord_message,
|
||||
format_file_size
|
||||
)
|
||||
from .discord.cdn_extractor import (
|
||||
extract_cdn_urls_from_response,
|
||||
send_cdn_urls_file
|
||||
)
|
||||
from .media.image_processing import tensor_to_numpy_uint8, process_batched_images
|
||||
from .github_integration import update_github_cdn_urls
|
||||
from .logging_config import setup_logging, get_logger
|
||||
from .filename_utils import build_filename_with_metadata, get_timestamp_string
|
||||
from .path_utils import get_output_directory, ensure_directory_exists, validate_path_is_safe
|
||||
|
||||
__all__ = [
|
||||
# Workflow utilities
|
||||
'sanitize_json_for_export',
|
||||
'extract_prompts_from_workflow',
|
||||
'WorkflowBuilder',
|
||||
# Discord utilities
|
||||
'DiscordWebhookClient',
|
||||
'validate_webhook_url',
|
||||
'send_to_discord_with_retry',
|
||||
'sanitize_token_from_text',
|
||||
# Discord message building
|
||||
'build_metadata_section',
|
||||
'build_prompt_section',
|
||||
'build_discord_message',
|
||||
'format_file_size',
|
||||
# CDN extraction
|
||||
'extract_cdn_urls_from_response',
|
||||
'send_cdn_urls_file',
|
||||
# Media utilities
|
||||
'tensor_to_numpy_uint8',
|
||||
'process_batched_images',
|
||||
# GitHub integration
|
||||
'update_github_cdn_urls',
|
||||
# Logging
|
||||
'setup_logging',
|
||||
'get_logger',
|
||||
# Filename utilities
|
||||
'build_filename_with_metadata',
|
||||
'get_timestamp_string',
|
||||
# Path utilities
|
||||
'get_output_directory',
|
||||
'ensure_directory_exists',
|
||||
'validate_path_is_safe',
|
||||
]
|
||||
@@ -0,0 +1,48 @@
|
||||
"""
|
||||
Discord Integration Utilities
|
||||
|
||||
Provides webhook client, message building, and CDN URL handling.
|
||||
"""
|
||||
|
||||
from .webhook_client import (
|
||||
DiscordWebhookClient,
|
||||
validate_webhook_url,
|
||||
sanitize_webhook_for_logging,
|
||||
send_to_discord_with_retry,
|
||||
validate_file_for_discord
|
||||
)
|
||||
from .message_builder import (
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
build_discord_message,
|
||||
validate_message_content,
|
||||
format_file_info,
|
||||
format_file_size
|
||||
)
|
||||
from .cdn_extractor import (
|
||||
extract_cdn_urls_from_response,
|
||||
create_cdn_urls_content,
|
||||
send_cdn_urls_file,
|
||||
collect_and_send_cdn_urls
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Webhook client
|
||||
'DiscordWebhookClient',
|
||||
'validate_webhook_url',
|
||||
'sanitize_webhook_for_logging',
|
||||
'send_to_discord_with_retry',
|
||||
'validate_file_for_discord',
|
||||
# Message building
|
||||
'build_metadata_section',
|
||||
'build_prompt_section',
|
||||
'build_discord_message',
|
||||
'validate_message_content',
|
||||
'format_file_info',
|
||||
'format_file_size',
|
||||
# CDN extraction
|
||||
'extract_cdn_urls_from_response',
|
||||
'create_cdn_urls_content',
|
||||
'send_cdn_urls_file',
|
||||
'collect_and_send_cdn_urls',
|
||||
]
|
||||
@@ -0,0 +1,168 @@
|
||||
"""
|
||||
Discord CDN URL extraction utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides functions for extracting CDN URLs from Discord responses
|
||||
and creating/sending URL text files.
|
||||
"""
|
||||
|
||||
from typing import List, Tuple, Optional, Any
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
def extract_cdn_urls_from_response(
|
||||
response: Any,
|
||||
exclude_json: bool = True
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Extract CDN URLs from a Discord webhook response.
|
||||
|
||||
Args:
|
||||
response: Response object from Discord API (must have .status_code and .json())
|
||||
exclude_json: Whether to exclude .json files from results
|
||||
|
||||
Returns:
|
||||
List of (filename, url) tuples
|
||||
"""
|
||||
cdn_urls = []
|
||||
|
||||
if response.status_code != 200:
|
||||
return cdn_urls
|
||||
|
||||
try:
|
||||
response_data = response.json()
|
||||
print(f"Received JSON response from Discord with "
|
||||
f"{len(response_data) if isinstance(response_data, dict) else 'invalid'} fields")
|
||||
|
||||
if "attachments" in response_data and isinstance(response_data["attachments"], list):
|
||||
print(f"Found {len(response_data['attachments'])} attachments in Discord response")
|
||||
|
||||
for idx, attachment in enumerate(response_data["attachments"]):
|
||||
if "url" in attachment and "filename" in attachment:
|
||||
filename = attachment["filename"]
|
||||
url = attachment["url"]
|
||||
|
||||
# Filter out workflow JSON files if requested
|
||||
if exclude_json and filename.endswith(".json"):
|
||||
print(f"Skipping JSON file: {filename}")
|
||||
continue
|
||||
|
||||
cdn_urls.append((filename, url))
|
||||
print(f"Extracted CDN URL for attachment {idx + 1}: {url}")
|
||||
else:
|
||||
print(f"Attachment {idx + 1} missing URL or filename: {attachment.keys()}")
|
||||
|
||||
print(f"Total CDN URLs collected: {len(cdn_urls)}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error extracting CDN URLs from response: {e}")
|
||||
|
||||
return cdn_urls
|
||||
|
||||
|
||||
def create_cdn_urls_content(
|
||||
urls: List[Tuple[str, str]],
|
||||
header: str = "# Discord CDN URLs\n\n"
|
||||
) -> str:
|
||||
"""
|
||||
Create text content from a list of CDN URLs.
|
||||
|
||||
Args:
|
||||
urls: List of (filename, url) tuples
|
||||
header: Header text for the content
|
||||
|
||||
Returns:
|
||||
Formatted text content
|
||||
"""
|
||||
content = header
|
||||
for idx, (filename, url) in enumerate(urls):
|
||||
content += f"{idx + 1}. {filename}: {url}\n"
|
||||
return content
|
||||
|
||||
|
||||
def send_cdn_urls_file(
|
||||
webhook_url: str,
|
||||
urls: List[Tuple[str, str]],
|
||||
send_func: Any,
|
||||
message: str = "Discord CDN URLs for the uploaded files:",
|
||||
filename_prefix: str = "cdn_urls"
|
||||
) -> bool:
|
||||
"""
|
||||
Create and send a text file containing CDN URLs to Discord.
|
||||
|
||||
Args:
|
||||
webhook_url: Discord webhook URL
|
||||
urls: List of (filename, url) tuples
|
||||
send_func: Function to send to Discord (send_to_discord_with_retry)
|
||||
message: Message to accompany the file
|
||||
filename_prefix: Prefix for the generated filename
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not urls:
|
||||
print("No CDN URLs to send")
|
||||
return False
|
||||
|
||||
try:
|
||||
# Create the text file content
|
||||
url_text_content = create_cdn_urls_content(urls)
|
||||
|
||||
# Create a unique filename for the text file
|
||||
urls_filename = f"{filename_prefix}-{uuid4()}.txt"
|
||||
|
||||
# Prepare the request with just the URL file
|
||||
url_files = {"file": (urls_filename, url_text_content.encode('utf-8'))}
|
||||
url_data = {"content": message}
|
||||
|
||||
# Send a follow-up message with just the URLs text file
|
||||
url_response = send_func(
|
||||
webhook_url,
|
||||
files=url_files,
|
||||
data=url_data
|
||||
)
|
||||
|
||||
if url_response.status_code in [200, 204]:
|
||||
print(f"Successfully sent CDN URLs text file to Discord")
|
||||
return True
|
||||
else:
|
||||
print(f"Error sending CDN URLs text file: Status code {url_response.status_code}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error creating or sending CDN URLs text file: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def collect_and_send_cdn_urls(
|
||||
response: Any,
|
||||
webhook_url: str,
|
||||
send_func: Any,
|
||||
save_cdn_urls: bool,
|
||||
existing_urls: Optional[List[Tuple[str, str]]] = None,
|
||||
message: str = "Discord CDN URLs for the uploaded files:"
|
||||
) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
Convenience function to extract CDN URLs from a response and optionally send them.
|
||||
|
||||
Args:
|
||||
response: Discord webhook response
|
||||
webhook_url: Webhook URL for sending the URLs file
|
||||
send_func: Function to send to Discord
|
||||
save_cdn_urls: Whether to extract and save CDN URLs
|
||||
existing_urls: Existing URLs to append to (for batch operations)
|
||||
message: Message to accompany the URLs file
|
||||
|
||||
Returns:
|
||||
List of all collected CDN URLs
|
||||
"""
|
||||
if existing_urls is None:
|
||||
existing_urls = []
|
||||
|
||||
if not save_cdn_urls:
|
||||
return existing_urls
|
||||
|
||||
# Extract URLs from this response
|
||||
new_urls = extract_cdn_urls_from_response(response)
|
||||
all_urls = existing_urls + new_urls
|
||||
|
||||
return all_urls
|
||||
@@ -0,0 +1,210 @@
|
||||
"""
|
||||
Discord message building utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides functions for constructing Discord messages with metadata,
|
||||
prompts, and other formatted content.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional, Tuple, List
|
||||
|
||||
|
||||
def build_metadata_section(
|
||||
info_dict: Dict[str, Any],
|
||||
include_date: bool = True,
|
||||
include_time: bool = True,
|
||||
include_dimensions: bool = True,
|
||||
include_format: bool = True,
|
||||
file_format: Optional[str] = None,
|
||||
frame_rate: Optional[float] = None,
|
||||
section_title: str = "Information"
|
||||
) -> str:
|
||||
"""
|
||||
Build a formatted metadata section for Discord messages.
|
||||
|
||||
Args:
|
||||
info_dict: Dictionary containing metadata (date, time, dimensions, etc.)
|
||||
include_date: Whether to include date if present
|
||||
include_time: Whether to include time if present
|
||||
include_dimensions: Whether to include dimensions if present
|
||||
include_format: Whether to include format information
|
||||
file_format: File format string (e.g., "png", "mp4")
|
||||
frame_rate: Frame rate for video (optional)
|
||||
section_title: Title for the section (e.g., "Image Information", "Video Info")
|
||||
|
||||
Returns:
|
||||
Formatted metadata string, or empty string if no metadata
|
||||
"""
|
||||
metadata_lines = []
|
||||
|
||||
if include_date and "date" in info_dict:
|
||||
metadata_lines.append(f"**Date:** {info_dict['date']}")
|
||||
|
||||
if include_time and "time" in info_dict:
|
||||
metadata_lines.append(f"**Time:** {info_dict['time']}")
|
||||
|
||||
if include_dimensions and "dimensions" in info_dict:
|
||||
metadata_lines.append(f"**Dimensions:** {info_dict['dimensions']}")
|
||||
|
||||
if frame_rate is not None:
|
||||
metadata_lines.append(f"**Frame Rate:** {frame_rate} fps")
|
||||
|
||||
if include_format and file_format:
|
||||
metadata_lines.append(f"**Format:** {file_format.upper()}")
|
||||
|
||||
if not metadata_lines:
|
||||
return ""
|
||||
|
||||
section = f"\n\n**{section_title}:**\n"
|
||||
section += "\n".join(metadata_lines) + "\n"
|
||||
return section
|
||||
|
||||
|
||||
def build_prompt_section(
|
||||
positive_prompt: Optional[str],
|
||||
negative_prompt: Optional[str],
|
||||
section_title: str = "Generation Prompts"
|
||||
) -> str:
|
||||
"""
|
||||
Build a formatted prompts section for Discord messages.
|
||||
|
||||
Args:
|
||||
positive_prompt: The positive/main prompt text
|
||||
negative_prompt: The negative prompt text
|
||||
section_title: Title for the section
|
||||
|
||||
Returns:
|
||||
Formatted prompt string, or empty string if no prompts
|
||||
"""
|
||||
# Validate and normalize prompts
|
||||
if positive_prompt is not None and not isinstance(positive_prompt, str):
|
||||
positive_prompt = str(positive_prompt)
|
||||
if negative_prompt is not None and not isinstance(negative_prompt, str):
|
||||
negative_prompt = str(negative_prompt)
|
||||
|
||||
has_positive = isinstance(positive_prompt, str) and positive_prompt.strip()
|
||||
has_negative = isinstance(negative_prompt, str) and negative_prompt.strip()
|
||||
|
||||
if not has_positive and not has_negative:
|
||||
return ""
|
||||
|
||||
section = f"\n\n**{section_title}:**\n"
|
||||
|
||||
if has_positive:
|
||||
section += f"**Positive:**\n```\n{positive_prompt.strip()}\n```\n"
|
||||
|
||||
if has_negative:
|
||||
section += f"**Negative:**\n```\n{negative_prompt.strip()}\n```\n"
|
||||
|
||||
return section
|
||||
|
||||
|
||||
def build_discord_message(
|
||||
base_message: str = "",
|
||||
metadata_section: str = "",
|
||||
prompt_section: str = "",
|
||||
additional_sections: Optional[List[str]] = None,
|
||||
max_length: int = 2000
|
||||
) -> str:
|
||||
"""
|
||||
Build a complete Discord message from components.
|
||||
|
||||
Args:
|
||||
base_message: The main message content
|
||||
metadata_section: Pre-built metadata section
|
||||
prompt_section: Pre-built prompt section
|
||||
additional_sections: List of additional section strings
|
||||
max_length: Maximum message length (Discord limit is 2000)
|
||||
|
||||
Returns:
|
||||
Complete formatted message, truncated if necessary
|
||||
"""
|
||||
parts = [base_message] if base_message else []
|
||||
|
||||
if metadata_section:
|
||||
parts.append(metadata_section)
|
||||
|
||||
if prompt_section:
|
||||
parts.append(prompt_section)
|
||||
|
||||
if additional_sections:
|
||||
parts.extend(additional_sections)
|
||||
|
||||
message = "".join(parts)
|
||||
|
||||
# Truncate if necessary
|
||||
if len(message) > max_length:
|
||||
truncation_notice = "\n...[Message truncated]"
|
||||
message = message[:max_length - len(truncation_notice)] + truncation_notice
|
||||
|
||||
return message
|
||||
|
||||
|
||||
def validate_message_content(message: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate Discord message content.
|
||||
|
||||
Args:
|
||||
message: Message content to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, validation_message)
|
||||
"""
|
||||
if not message:
|
||||
return True, "Empty message (valid for file-only uploads)"
|
||||
|
||||
if len(message) > 2000:
|
||||
return False, f"Message exceeds 2000 character limit ({len(message)} chars)"
|
||||
|
||||
# Check for required sections (informational)
|
||||
has_prompts = "Generation Prompts" in message
|
||||
|
||||
info_parts = []
|
||||
info_parts.append(f"Message has {message.count(chr(10))} lines")
|
||||
|
||||
if not has_prompts:
|
||||
info_parts.append("WARNING: Message does NOT contain 'Generation Prompts' section")
|
||||
|
||||
return True, "\n".join(info_parts)
|
||||
|
||||
|
||||
def format_file_info(
|
||||
filename: str,
|
||||
file_size: int,
|
||||
mime_type: Optional[str] = None
|
||||
) -> str:
|
||||
"""
|
||||
Format file information for logging/display.
|
||||
|
||||
Args:
|
||||
filename: Name of the file
|
||||
file_size: Size in bytes
|
||||
mime_type: MIME type of the file
|
||||
|
||||
Returns:
|
||||
Formatted string with file information
|
||||
"""
|
||||
size_str = format_file_size(file_size)
|
||||
info = f"File: {filename} ({size_str})"
|
||||
if mime_type:
|
||||
info += f" [{mime_type}]"
|
||||
return info
|
||||
|
||||
|
||||
def format_file_size(size_bytes: int) -> str:
|
||||
"""
|
||||
Format file size in human-readable format.
|
||||
|
||||
Args:
|
||||
size_bytes: Size in bytes
|
||||
|
||||
Returns:
|
||||
Formatted string (e.g., "1.5 MB", "256 KB")
|
||||
"""
|
||||
if size_bytes < 1024:
|
||||
return f"{size_bytes} bytes"
|
||||
elif size_bytes < 1024 * 1024:
|
||||
return f"{size_bytes / 1024:.1f} KB"
|
||||
elif size_bytes < 1024 * 1024 * 1024:
|
||||
return f"{size_bytes / (1024 * 1024):.1f} MB"
|
||||
else:
|
||||
return f"{size_bytes / (1024 * 1024 * 1024):.2f} GB"
|
||||
@@ -0,0 +1,458 @@
|
||||
"""
|
||||
Discord API Utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides a client for interacting with Discord webhooks and validation utilities.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import logging
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import json
|
||||
import requests
|
||||
|
||||
# Get logger for this module
|
||||
logger = logging.getLogger("comfyui_discordsend")
|
||||
|
||||
|
||||
# Discord webhook URL patterns
|
||||
WEBHOOK_URL_PATTERNS = [
|
||||
r"https://(?:www\.)?discord(?:app)?\.com/api/webhooks/\d+/[\w-]+$",
|
||||
r"https://(?:www\.)?discordapp\.com/api/webhooks/\d+/[\w-]+$",
|
||||
]
|
||||
|
||||
|
||||
def validate_webhook_url(url: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate a Discord webhook URL.
|
||||
|
||||
Args:
|
||||
url: The webhook URL to validate
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, message)
|
||||
"""
|
||||
if not url:
|
||||
return False, "Webhook URL is empty"
|
||||
|
||||
if not url.startswith("https://"):
|
||||
return False, "Webhook URL must start with https://"
|
||||
|
||||
# Check against known patterns
|
||||
for pattern in WEBHOOK_URL_PATTERNS:
|
||||
if re.match(pattern, url, re.IGNORECASE):
|
||||
return True, "Valid Discord webhook URL"
|
||||
|
||||
return False, "URL does not appear to be a valid Discord webhook URL"
|
||||
|
||||
|
||||
def sanitize_webhook_for_logging(url: str) -> str:
|
||||
"""
|
||||
Sanitize a webhook URL for safe logging (hide the token portion).
|
||||
|
||||
Args:
|
||||
url: The webhook URL
|
||||
|
||||
Returns:
|
||||
Sanitized URL safe for logging
|
||||
"""
|
||||
if not url:
|
||||
return ""
|
||||
|
||||
# Pattern: https://discord.com/api/webhooks/{id}/{token}
|
||||
match = re.match(r"(https?://[^/]+/api/webhooks/\d+/)(.+)", url)
|
||||
if match:
|
||||
return f"{match.group(1)}[REDACTED]"
|
||||
|
||||
return "[REDACTED_WEBHOOK_URL]"
|
||||
|
||||
|
||||
def sanitize_token_from_text(text: str, webhook_url: str) -> str:
|
||||
"""
|
||||
Sanitize the webhook token from arbitrary text.
|
||||
|
||||
Args:
|
||||
text: The text to sanitize
|
||||
webhook_url: The webhook URL containing the token
|
||||
|
||||
Returns:
|
||||
Text with the token replaced by [REDACTED]
|
||||
"""
|
||||
if not text or not webhook_url:
|
||||
return text
|
||||
|
||||
# Pattern: https://discord.com/api/webhooks/{id}/{token}
|
||||
# Use case-insensitive matching to handle potential variations
|
||||
match = re.search(r"/api/webhooks/\d+/([\w-]+)", webhook_url, re.IGNORECASE)
|
||||
if match:
|
||||
token = match.group(1)
|
||||
if token in text:
|
||||
return text.replace(token, "[REDACTED]")
|
||||
|
||||
return text
|
||||
|
||||
|
||||
class DiscordWebhookClient:
|
||||
"""
|
||||
Client for sending messages and files to Discord via webhooks.
|
||||
|
||||
Features:
|
||||
- Automatic retry with exponential backoff
|
||||
- Rate limit handling
|
||||
- File size validation
|
||||
- Error handling with sanitized logging
|
||||
"""
|
||||
|
||||
# Discord limits
|
||||
MAX_FILE_SIZE = 25 * 1024 * 1024 # 25MB for regular users
|
||||
MAX_FILES_PER_MESSAGE = 10
|
||||
MAX_MESSAGE_LENGTH = 2000
|
||||
|
||||
def __init__(self, webhook_url: str, max_retries: int = 3):
|
||||
"""
|
||||
Initialize the Discord webhook client.
|
||||
|
||||
Args:
|
||||
webhook_url: The Discord webhook URL
|
||||
max_retries: Maximum number of retry attempts for failed requests
|
||||
"""
|
||||
self.webhook_url = webhook_url
|
||||
self.max_retries = max_retries
|
||||
self._validated = False
|
||||
|
||||
def validate(self) -> Tuple[bool, str]:
|
||||
"""Validate the webhook URL."""
|
||||
is_valid, message = validate_webhook_url(self.webhook_url)
|
||||
self._validated = is_valid
|
||||
return is_valid, message
|
||||
|
||||
def send_message(
|
||||
self,
|
||||
content: str = "",
|
||||
files: Optional[List[Tuple[str, bytes, str]]] = None,
|
||||
embeds: Optional[List[Dict]] = None
|
||||
) -> Tuple[bool, Dict[str, Any]]:
|
||||
"""
|
||||
Send a message to Discord.
|
||||
|
||||
Args:
|
||||
content: Text message content
|
||||
files: List of (filename, file_bytes, content_type) tuples
|
||||
embeds: List of embed dictionaries
|
||||
|
||||
Returns:
|
||||
Tuple of (success, response_data)
|
||||
"""
|
||||
if not self.webhook_url:
|
||||
return False, {"error": "No webhook URL configured"}
|
||||
|
||||
# Validate URL if not already done
|
||||
if not self._validated:
|
||||
is_valid, message = self.validate()
|
||||
if not is_valid:
|
||||
return False, {"error": message}
|
||||
|
||||
# Truncate message if too long
|
||||
if content and len(content) > self.MAX_MESSAGE_LENGTH:
|
||||
content = content[:self.MAX_MESSAGE_LENGTH - 3] + "..."
|
||||
|
||||
# Prepare request
|
||||
data = {}
|
||||
if content:
|
||||
data["content"] = content
|
||||
if embeds:
|
||||
data["embeds"] = embeds
|
||||
|
||||
# Prepare files
|
||||
request_files = None
|
||||
if files:
|
||||
request_files = []
|
||||
for i, (filename, file_bytes, content_type) in enumerate(files):
|
||||
# Validate file size
|
||||
if len(file_bytes) > self.MAX_FILE_SIZE:
|
||||
continue # Skip oversized files
|
||||
|
||||
request_files.append(
|
||||
(f"file{i}", (filename, BytesIO(file_bytes), content_type))
|
||||
)
|
||||
|
||||
if len(request_files) > self.MAX_FILES_PER_MESSAGE:
|
||||
request_files = request_files[:self.MAX_FILES_PER_MESSAGE]
|
||||
|
||||
# Send with retry logic
|
||||
return self._send_with_retry(data, request_files)
|
||||
|
||||
def send_file(
|
||||
self,
|
||||
file_path: str,
|
||||
message: str = "",
|
||||
additional_files: Optional[List[str]] = None
|
||||
) -> Tuple[bool, Dict[str, Any]]:
|
||||
"""
|
||||
Send a file from disk to Discord.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file to send
|
||||
message: Optional message to accompany the file
|
||||
additional_files: Optional list of additional file paths
|
||||
|
||||
Returns:
|
||||
Tuple of (success, response_data)
|
||||
"""
|
||||
files = []
|
||||
|
||||
# Add main file
|
||||
if os.path.exists(file_path):
|
||||
file_size = os.path.getsize(file_path)
|
||||
if file_size <= self.MAX_FILE_SIZE:
|
||||
with open(file_path, "rb") as f:
|
||||
filename = os.path.basename(file_path)
|
||||
content_type = self._get_content_type(filename)
|
||||
files.append((filename, f.read(), content_type))
|
||||
|
||||
# Add additional files
|
||||
if additional_files:
|
||||
for path in additional_files:
|
||||
if os.path.exists(path):
|
||||
file_size = os.path.getsize(path)
|
||||
if file_size <= self.MAX_FILE_SIZE:
|
||||
with open(path, "rb") as f:
|
||||
filename = os.path.basename(path)
|
||||
content_type = self._get_content_type(filename)
|
||||
files.append((filename, f.read(), content_type))
|
||||
|
||||
if not files:
|
||||
return False, {"error": "No valid files to send"}
|
||||
|
||||
return self.send_message(content=message, files=files)
|
||||
|
||||
def _send_with_retry(
|
||||
self,
|
||||
data: Dict,
|
||||
files: Optional[List] = None
|
||||
) -> Tuple[bool, Dict[str, Any]]:
|
||||
"""Send request with retry logic."""
|
||||
last_error = None
|
||||
|
||||
for attempt in range(self.max_retries):
|
||||
try:
|
||||
if files:
|
||||
response = requests.post(
|
||||
self.webhook_url,
|
||||
data={"payload_json": json.dumps(data)} if data else None,
|
||||
files=files,
|
||||
timeout=60
|
||||
)
|
||||
else:
|
||||
response = requests.post(
|
||||
self.webhook_url,
|
||||
json=data,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
# Handle rate limiting
|
||||
if response.status_code == 429:
|
||||
retry_after = response.json().get("retry_after", 1)
|
||||
time.sleep(retry_after)
|
||||
continue
|
||||
|
||||
# Success
|
||||
if response.status_code in [200, 204]:
|
||||
try:
|
||||
return True, response.json() if response.content else {}
|
||||
except:
|
||||
return True, {}
|
||||
|
||||
# Client errors (don't retry)
|
||||
if 400 <= response.status_code < 500:
|
||||
sanitized_details = sanitize_token_from_text(response.text[:500], self.webhook_url)
|
||||
return False, {
|
||||
"error": f"Discord API error: {response.status_code}",
|
||||
"details": sanitized_details
|
||||
}
|
||||
|
||||
# Server errors (retry)
|
||||
last_error = f"Discord API returned {response.status_code}"
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
last_error = "Request timed out"
|
||||
except requests.exceptions.RequestException as e:
|
||||
# Sanitize error message to prevent token leakage
|
||||
error_msg = sanitize_token_from_text(str(e), self.webhook_url)
|
||||
last_error = error_msg
|
||||
|
||||
# Exponential backoff
|
||||
if attempt < self.max_retries - 1:
|
||||
time.sleep(2 ** attempt)
|
||||
|
||||
return False, {"error": last_error or "Unknown error"}
|
||||
|
||||
def _get_content_type(self, filename: str) -> str:
|
||||
"""Get content type based on file extension."""
|
||||
ext = os.path.splitext(filename)[1].lower()
|
||||
content_types = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".mp4": "video/mp4",
|
||||
".webm": "video/webm",
|
||||
".mov": "video/quicktime",
|
||||
".json": "application/json",
|
||||
".txt": "text/plain",
|
||||
}
|
||||
return content_types.get(ext, "application/octet-stream")
|
||||
|
||||
|
||||
def validate_file_for_discord(file_path: str) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate that a file is compatible with Discord uploads.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, message)
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
return False, f"File does not exist: {file_path}"
|
||||
|
||||
file_size = os.path.getsize(file_path)
|
||||
|
||||
if file_size == 0:
|
||||
return False, "File is empty"
|
||||
|
||||
if file_size < 1024:
|
||||
return False, f"File is suspiciously small: {file_size} bytes"
|
||||
|
||||
max_size = 25 * 1024 * 1024 # 25MB
|
||||
if file_size > max_size:
|
||||
return False, f"File exceeds Discord's size limit: {file_size} bytes (max {max_size} bytes)"
|
||||
|
||||
# Check extension
|
||||
ext = os.path.splitext(file_path)[1].lower().lstrip('.')
|
||||
supported_formats = ['png', 'jpg', 'jpeg', 'gif', 'webp', 'mp4', 'webm', 'mov']
|
||||
|
||||
if ext in supported_formats:
|
||||
return True, f"Valid {ext.upper()} file for Discord"
|
||||
else:
|
||||
return False, f"Format {ext} may not be fully supported by Discord"
|
||||
|
||||
|
||||
def send_to_discord_with_retry(
|
||||
webhook_url: str,
|
||||
files: Optional[List] = None,
|
||||
data: Optional[Dict] = None,
|
||||
json_data: Optional[Dict] = None,
|
||||
max_retries: int = 3,
|
||||
timeout: int = 60
|
||||
) -> requests.Response:
|
||||
"""
|
||||
Send a request to Discord webhook with retry logic.
|
||||
|
||||
This is a drop-in replacement for requests.post() with added retry logic,
|
||||
rate limit handling, and exponential backoff.
|
||||
|
||||
Args:
|
||||
webhook_url: The Discord webhook URL
|
||||
files: Files to upload (same format as requests.post)
|
||||
data: Form data (same format as requests.post)
|
||||
json_data: JSON data (same format as requests.post)
|
||||
max_retries: Maximum number of retry attempts
|
||||
timeout: Request timeout in seconds
|
||||
|
||||
Returns:
|
||||
The response object from the successful request
|
||||
|
||||
Raises:
|
||||
requests.exceptions.RequestException: If all retries fail
|
||||
ValueError: If the webhook URL is invalid
|
||||
"""
|
||||
# Validate URL to prevent SSRF
|
||||
is_valid, error_msg = validate_webhook_url(webhook_url)
|
||||
if not is_valid:
|
||||
raise ValueError(f"Invalid webhook URL: {error_msg}")
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
if files:
|
||||
response = requests.post(
|
||||
webhook_url,
|
||||
files=files,
|
||||
data=data,
|
||||
timeout=timeout
|
||||
)
|
||||
elif json_data:
|
||||
response = requests.post(
|
||||
webhook_url,
|
||||
json=json_data,
|
||||
timeout=timeout
|
||||
)
|
||||
else:
|
||||
response = requests.post(
|
||||
webhook_url,
|
||||
data=data,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# Handle rate limiting
|
||||
if response.status_code == 429:
|
||||
retry_after = 1
|
||||
try:
|
||||
retry_after = response.json().get("retry_after", 1)
|
||||
except:
|
||||
pass
|
||||
logger.warning(f"Rate limited by Discord, waiting {retry_after}s before retry...")
|
||||
time.sleep(retry_after)
|
||||
continue
|
||||
|
||||
# Success or client error (don't retry client errors)
|
||||
if response.status_code < 500:
|
||||
return response
|
||||
|
||||
# Server error - retry
|
||||
logger.warning(f"Discord server error {response.status_code}, attempt {attempt + 1}/{max_retries}")
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
logger.warning(f"Request timeout, attempt {attempt + 1}/{max_retries}")
|
||||
last_exception = requests.exceptions.Timeout("Discord request timed out")
|
||||
except requests.exceptions.RequestException as e:
|
||||
# Sanitize error message to prevent token leakage
|
||||
error_msg = sanitize_token_from_text(str(e), webhook_url)
|
||||
logger.warning(f"Request error: {error_msg}, attempt {attempt + 1}/{max_retries}")
|
||||
|
||||
# Store sanitized exception to avoid leaking token if raised later
|
||||
# Use case-insensitive matching to handle uppercase URLs
|
||||
match = re.search(r"/api/webhooks/\d+/([\w-]+)", webhook_url, re.IGNORECASE)
|
||||
if match and match.group(1) in str(e):
|
||||
# Create a new exception of the same type with sanitized message
|
||||
# We try to preserve the exception type, but fallback to RequestException if init fails
|
||||
try:
|
||||
last_exception = type(e)(error_msg)
|
||||
# Preserve context attributes if possible
|
||||
last_exception.request = getattr(e, "request", None)
|
||||
last_exception.response = getattr(e, "response", None)
|
||||
except:
|
||||
last_exception = requests.exceptions.RequestException(error_msg)
|
||||
last_exception.request = getattr(e, "request", None)
|
||||
last_exception.response = getattr(e, "response", None)
|
||||
else:
|
||||
last_exception = e
|
||||
|
||||
# Exponential backoff before retry
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = 2 ** attempt
|
||||
time.sleep(wait_time)
|
||||
|
||||
# If we get here, all retries failed
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
|
||||
# Return the last response even if it was an error
|
||||
return response
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
Filename utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides functions for building filenames with date, time, and dimension metadata.
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Dict, Optional, Tuple, Any
|
||||
|
||||
|
||||
def build_filename_with_metadata(
|
||||
prefix: str,
|
||||
add_date: bool = False,
|
||||
add_time: bool = False,
|
||||
add_dimensions: bool = False,
|
||||
width: Optional[int] = None,
|
||||
height: Optional[int] = None,
|
||||
info_dict: Optional[Dict[str, Any]] = None
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Build a filename with optional date, time, and dimension suffixes.
|
||||
|
||||
Args:
|
||||
prefix: The base filename prefix
|
||||
add_date: Whether to add the current date (YYYY-MM-DD)
|
||||
add_time: Whether to add the current time (HH-MM-SS)
|
||||
add_dimensions: Whether to add dimensions (WxH)
|
||||
width: Image/video width (required if add_dimensions is True)
|
||||
height: Image/video height (required if add_dimensions is True)
|
||||
info_dict: Optional dict to update with metadata (creates new if None)
|
||||
|
||||
Returns:
|
||||
Tuple of (modified_prefix, info_dict with metadata)
|
||||
"""
|
||||
if info_dict is None:
|
||||
info_dict = {}
|
||||
|
||||
metadata_parts = []
|
||||
|
||||
if add_date:
|
||||
current_date = time.strftime("%Y-%m-%d")
|
||||
metadata_parts.append(current_date)
|
||||
info_dict["date"] = current_date
|
||||
print(f"Adding date to filename: {current_date}")
|
||||
|
||||
if add_time:
|
||||
current_time = time.strftime("%H-%M-%S")
|
||||
metadata_parts.append(current_time)
|
||||
info_dict["time"] = current_time
|
||||
print(f"Adding time to filename: {current_time}")
|
||||
|
||||
if add_dimensions and width is not None and height is not None:
|
||||
dim_text = f"{width}x{height}"
|
||||
metadata_parts.append(dim_text)
|
||||
info_dict["dimensions"] = dim_text
|
||||
print(f"Adding dimensions to filename: {dim_text}")
|
||||
|
||||
modified_prefix = prefix
|
||||
if metadata_parts:
|
||||
metadata_suffix = "_" + "_".join(metadata_parts)
|
||||
modified_prefix += metadata_suffix
|
||||
print(f"Final metadata suffix: {metadata_suffix}")
|
||||
|
||||
return modified_prefix, info_dict
|
||||
|
||||
|
||||
def get_timestamp_string(include_date: bool = True, include_time: bool = True) -> str:
|
||||
"""
|
||||
Get a formatted timestamp string.
|
||||
|
||||
Args:
|
||||
include_date: Include date in format YYYY-MM-DD
|
||||
include_time: Include time in format HH-MM-SS
|
||||
|
||||
Returns:
|
||||
Formatted timestamp string
|
||||
"""
|
||||
parts = []
|
||||
if include_date:
|
||||
parts.append(time.strftime("%Y-%m-%d"))
|
||||
if include_time:
|
||||
parts.append(time.strftime("%H-%M-%S"))
|
||||
return "_".join(parts) if parts else ""
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
GitHub Integration Utilities for ComfyUI-DiscordSend
|
||||
|
||||
Handles updating GitHub repositories with Discord CDN URLs.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import time
|
||||
import re
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def validate_github_repo(repo: str) -> bool:
|
||||
"""
|
||||
Validate GitHub repository format (username/repo).
|
||||
Strictly enforces alphanumeric, hyphens, underscores, and periods.
|
||||
Prevents path traversal and injection.
|
||||
"""
|
||||
if not repo:
|
||||
return False
|
||||
# Pattern: username/repo
|
||||
# GitHub usernames: alphanumeric, hyphens (max 39 chars)
|
||||
# Repo names: alphanumeric, hyphens, periods, underscores
|
||||
pattern = r"^[a-zA-Z0-9-]+/[\w.-]+$"
|
||||
return bool(re.match(pattern, repo))
|
||||
|
||||
|
||||
def validate_file_path(path: str) -> bool:
|
||||
"""
|
||||
Validate file path for GitHub API.
|
||||
Prevents path traversal (..) and absolute paths.
|
||||
"""
|
||||
if not path:
|
||||
return False
|
||||
# Prevent traversal
|
||||
if ".." in path:
|
||||
return False
|
||||
# Prevent absolute paths (GitHub API treats paths as relative to root)
|
||||
if path.startswith("/"):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def update_github_cdn_urls(
|
||||
github_repo: str,
|
||||
github_token: str,
|
||||
file_path: str,
|
||||
cdn_urls: List[Tuple[str, str]],
|
||||
commit_message: Optional[str] = None
|
||||
) -> Tuple[bool, str]:
|
||||
"""
|
||||
Update a file in a GitHub repository with Discord CDN URLs.
|
||||
|
||||
Args:
|
||||
github_repo: The GitHub repository (format: username/repo)
|
||||
github_token: The GitHub personal access token for authentication
|
||||
file_path: The path to the file within the repository to update
|
||||
cdn_urls: List of (filename, url) tuples containing Discord CDN URLs
|
||||
commit_message: Optional commit message, defaults to a standard message
|
||||
|
||||
Returns:
|
||||
Tuple of (success, message) where success is a boolean and message is a status message
|
||||
"""
|
||||
# Validate required parameters
|
||||
if not github_repo:
|
||||
return False, "Missing GitHub repository name"
|
||||
|
||||
if not github_token:
|
||||
return False, "Missing GitHub personal access token"
|
||||
|
||||
if not file_path:
|
||||
return False, "Missing file path in repository"
|
||||
|
||||
if not cdn_urls:
|
||||
return False, "No CDN URLs to update"
|
||||
|
||||
# Strictly validate repository format to prevent traversal/injection
|
||||
if not validate_github_repo(github_repo):
|
||||
return False, f"Invalid GitHub repository format: {github_repo}. Expected format: username/repo (alphanumeric, hyphens, periods, underscores only)"
|
||||
|
||||
# Strictly validate file path to prevent traversal
|
||||
if not validate_file_path(file_path):
|
||||
return False, f"Invalid file path: {file_path}. Path traversal (..) and absolute paths are not allowed."
|
||||
|
||||
# Setup API endpoint
|
||||
api_url = f"https://api.github.com/repos/{github_repo}/contents/{file_path}"
|
||||
|
||||
headers = {
|
||||
"Authorization": f"token {github_token}",
|
||||
"Accept": "application/vnd.github.v3+json"
|
||||
}
|
||||
|
||||
try:
|
||||
# Check if file exists and get its SHA
|
||||
file_sha = None
|
||||
current_content = ""
|
||||
|
||||
response = requests.get(api_url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
file_data = response.json()
|
||||
file_sha = file_data.get("sha")
|
||||
|
||||
# Get current content
|
||||
if file_data.get("content"):
|
||||
current_content = base64.b64decode(file_data["content"]).decode("utf-8")
|
||||
|
||||
elif response.status_code == 404:
|
||||
pass # File doesn't exist, will create new
|
||||
else:
|
||||
# Sanitize response text
|
||||
error_details = response.text
|
||||
if github_token and github_token in error_details:
|
||||
error_details = error_details.replace(github_token, "[REDACTED_TOKEN]")
|
||||
return False, f"Error checking GitHub file: {response.status_code} - {error_details}"
|
||||
|
||||
# Prepare file content
|
||||
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
# Merge with existing content if present
|
||||
existing_urls = {}
|
||||
if current_content:
|
||||
for line in current_content.splitlines():
|
||||
if ": https://" in line and "cdn.discordapp.com" in line:
|
||||
parts = line.split(": ", 1)
|
||||
if len(parts) == 2:
|
||||
name_part = parts[0]
|
||||
# Remove numbering if present
|
||||
if ". " in name_part:
|
||||
name_part = name_part.split(". ", 1)[1]
|
||||
existing_urls[name_part] = parts[1]
|
||||
|
||||
# Add new URLs (overwrites duplicates)
|
||||
for filename, url in cdn_urls:
|
||||
existing_urls[filename] = url
|
||||
|
||||
# Format content
|
||||
new_content = f"# Discord CDN URLs\nLast updated: {timestamp}\n\n"
|
||||
for i, (filename, url) in enumerate(existing_urls.items(), 1):
|
||||
new_content += f"{i}. {filename}: {url}\n"
|
||||
|
||||
# Set commit message
|
||||
if not commit_message:
|
||||
commit_message = f"Update Discord CDN URLs - {timestamp}"
|
||||
|
||||
# Prepare request data
|
||||
data = {
|
||||
"message": commit_message,
|
||||
"content": base64.b64encode(new_content.encode("utf-8")).decode("utf-8"),
|
||||
}
|
||||
|
||||
if file_sha:
|
||||
data["sha"] = file_sha
|
||||
|
||||
# Create/update file
|
||||
response = requests.put(api_url, headers=headers, json=data, timeout=30)
|
||||
|
||||
if response.status_code in [200, 201]:
|
||||
return True, f"Successfully updated GitHub file with {len(cdn_urls)} Discord CDN URLs"
|
||||
else:
|
||||
# Sanitize response text to ensure no token leakage
|
||||
error_details = response.text
|
||||
if github_token and github_token in error_details:
|
||||
error_details = error_details.replace(github_token, "[REDACTED_TOKEN]")
|
||||
return False, f"Error updating GitHub file: {response.status_code} - {error_details}"
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
return False, "GitHub API request timed out"
|
||||
except requests.exceptions.RequestException as e:
|
||||
# Scrub token from error messages
|
||||
error_message = str(e)
|
||||
if github_token and github_token in error_message:
|
||||
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
|
||||
return False, f"GitHub API request failed: {error_message}"
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
if github_token and github_token in error_message:
|
||||
error_message = error_message.replace(github_token, "[REDACTED_TOKEN]")
|
||||
return False, f"Exception during GitHub update: {error_message}"
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Logging configuration for ComfyUI-DiscordSend
|
||||
|
||||
Provides a configured logger for the extension.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
|
||||
|
||||
def setup_logging(level: int = logging.INFO) -> None:
|
||||
"""
|
||||
Set up logging configuration for the application.
|
||||
|
||||
Args:
|
||||
level: The logging level to use (default: INFO)
|
||||
"""
|
||||
logging.basicConfig(
|
||||
level=level,
|
||||
format='[%(name)s] %(levelname)s: %(message)s',
|
||||
handlers=[
|
||||
logging.StreamHandler(sys.stdout)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def get_logger(name: str = "comfyui_discordsend") -> logging.Logger:
|
||||
"""
|
||||
Get a configured logger for the extension.
|
||||
|
||||
Args:
|
||||
name: The logger name (default: comfyui_discordsend)
|
||||
|
||||
Returns:
|
||||
Configured logger instance
|
||||
"""
|
||||
logger = logging.getLogger(name)
|
||||
|
||||
# Only configure if not already configured
|
||||
if not logger.handlers:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
# Console handler
|
||||
console_handler = logging.StreamHandler(sys.stdout)
|
||||
console_handler.setLevel(logging.INFO)
|
||||
|
||||
# Format
|
||||
formatter = logging.Formatter(
|
||||
'[%(name)s] %(levelname)s: %(message)s'
|
||||
)
|
||||
console_handler.setFormatter(formatter)
|
||||
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
# Default logger instance
|
||||
logger = get_logger()
|
||||
@@ -0,0 +1,40 @@
|
||||
"""
|
||||
Media Processing Utilities
|
||||
|
||||
Provides image and video processing functions.
|
||||
"""
|
||||
|
||||
from .image_processing import tensor_to_numpy_uint8
|
||||
from .format_utils import (
|
||||
parse_format_string,
|
||||
normalize_video_extension,
|
||||
get_mime_type,
|
||||
validate_video_for_discord,
|
||||
is_animated_format,
|
||||
supports_alpha
|
||||
)
|
||||
from .video_encoder import (
|
||||
detect_ffmpeg,
|
||||
FFmpegEncoder,
|
||||
PILEncoder,
|
||||
optimize_video_for_discord,
|
||||
mux_audio_to_video
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Image processing
|
||||
'tensor_to_numpy_uint8',
|
||||
# Format utilities
|
||||
'parse_format_string',
|
||||
'normalize_video_extension',
|
||||
'get_mime_type',
|
||||
'validate_video_for_discord',
|
||||
'is_animated_format',
|
||||
'supports_alpha',
|
||||
# Video encoding
|
||||
'detect_ffmpeg',
|
||||
'FFmpegEncoder',
|
||||
'PILEncoder',
|
||||
'optimize_video_for_discord',
|
||||
'mux_audio_to_video',
|
||||
]
|
||||
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
Video format utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides format detection, extension mapping, and validation.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Tuple, Optional
|
||||
|
||||
|
||||
def parse_format_string(format_str: str) -> Tuple[str, str]:
|
||||
"""
|
||||
Parse a format string into type and extension.
|
||||
|
||||
Args:
|
||||
format_str: Format string like "video/h264-mp4" or "image/gif"
|
||||
|
||||
Returns:
|
||||
Tuple of (format_type, format_extension)
|
||||
"""
|
||||
if "/" in format_str:
|
||||
format_type, format_ext = format_str.split("/", 1)
|
||||
else:
|
||||
format_type = "video"
|
||||
format_ext = format_str
|
||||
|
||||
return format_type, format_ext
|
||||
|
||||
|
||||
def normalize_video_extension(format_str: str) -> str:
|
||||
"""
|
||||
Normalize a format string to a file extension.
|
||||
|
||||
Args:
|
||||
format_str: Format string like "video/h264-mp4"
|
||||
|
||||
Returns:
|
||||
Normalized extension (e.g., "mp4", "webm", "gif")
|
||||
"""
|
||||
_, format_ext = parse_format_string(format_str)
|
||||
|
||||
# Map format strings to extensions
|
||||
extension_map = {
|
||||
"h264-mp4": "mp4",
|
||||
"h265-mp4": "mp4",
|
||||
"vp9-webm": "webm",
|
||||
"prores": "mov",
|
||||
}
|
||||
|
||||
return extension_map.get(format_ext, format_ext)
|
||||
|
||||
|
||||
def get_mime_type(extension: str) -> str:
|
||||
"""
|
||||
Get MIME type for a video extension.
|
||||
|
||||
Args:
|
||||
extension: File extension (without dot)
|
||||
|
||||
Returns:
|
||||
MIME type string
|
||||
"""
|
||||
mime_types = {
|
||||
"mp4": "video/mp4",
|
||||
"webm": "video/webm",
|
||||
"gif": "image/gif",
|
||||
"mov": "video/quicktime",
|
||||
"avi": "video/x-msvideo",
|
||||
"mkv": "video/x-matroska",
|
||||
}
|
||||
return mime_types.get(extension.lower(), "application/octet-stream")
|
||||
|
||||
|
||||
def validate_video_for_discord(file_path: str, max_size_mb: int = 25) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate that a video file is compatible with Discord.
|
||||
|
||||
Args:
|
||||
file_path: Path to the video file
|
||||
max_size_mb: Maximum file size in megabytes (default 25MB for Discord)
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, message)
|
||||
"""
|
||||
if not os.path.exists(file_path):
|
||||
return False, f"File does not exist: {file_path}"
|
||||
|
||||
file_size = os.path.getsize(file_path)
|
||||
|
||||
if file_size == 0:
|
||||
return False, "File is empty"
|
||||
|
||||
if file_size < 1024:
|
||||
return False, f"File is suspiciously small: {file_size} bytes"
|
||||
|
||||
max_size = max_size_mb * 1024 * 1024
|
||||
if file_size > max_size:
|
||||
return False, f"File exceeds Discord's size limit of {max_size_mb}MB ({file_size / (1024*1024):.2f}MB)"
|
||||
|
||||
ext = os.path.splitext(file_path)[1].lower().lstrip('.')
|
||||
|
||||
if ext in ['mp4', 'webm', 'gif']:
|
||||
return True, "Valid"
|
||||
elif ext in ['mov']:
|
||||
return False, "MOV files may need conversion for Discord compatibility"
|
||||
elif ext in ['png', 'apng']:
|
||||
return False, "PNG/APNG sequence may need compilation for Discord"
|
||||
else:
|
||||
return False, f"Unknown format '{ext}' - may not be compatible with Discord"
|
||||
|
||||
|
||||
def is_animated_format(extension: str) -> bool:
|
||||
"""
|
||||
Check if a format supports animation.
|
||||
|
||||
Args:
|
||||
extension: File extension (without dot)
|
||||
|
||||
Returns:
|
||||
True if the format supports animation
|
||||
"""
|
||||
animated_formats = {'gif', 'webp', 'mp4', 'webm', 'mov', 'avi', 'mkv', 'apng'}
|
||||
return extension.lower() in animated_formats
|
||||
|
||||
|
||||
def supports_alpha(extension: str) -> bool:
|
||||
"""
|
||||
Check if a format supports alpha channel (transparency).
|
||||
|
||||
Args:
|
||||
extension: File extension (without dot)
|
||||
|
||||
Returns:
|
||||
True if the format supports alpha
|
||||
"""
|
||||
alpha_formats = {'webm', 'gif', 'webp', 'png', 'apng', 'mov'}
|
||||
return extension.lower() in alpha_formats
|
||||
@@ -0,0 +1,51 @@
|
||||
"""
|
||||
Image processing utilities for ComfyUI-DiscordSend.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
def tensor_to_numpy_uint8(tensor: torch.Tensor) -> np.ndarray:
|
||||
"""
|
||||
Convert a PyTorch tensor (0-1 float) to a numpy uint8 array (0-255).
|
||||
|
||||
This function optimizes performance by doing scaling, clamping, and casting
|
||||
in PyTorch before moving data to CPU/NumPy, avoiding large intermediate float arrays.
|
||||
|
||||
Args:
|
||||
tensor: PyTorch tensor with values in range [0, 1]
|
||||
|
||||
Returns:
|
||||
Numpy uint8 array with values in range [0, 255]
|
||||
"""
|
||||
# Optimization: Use torch operations for scaling/clipping/casting to avoid large float64 intermediate arrays on CPU
|
||||
# This is ~70% faster than naive numpy conversion: np.clip(255. * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
||||
# Further Optimization: Use clamp_ (in-place) to avoid allocating a second float tensor
|
||||
return (tensor * 255.0).clamp_(0, 255).to(dtype=torch.uint8).cpu().numpy()
|
||||
|
||||
def process_batched_images(image_sequence, batch_size=20):
|
||||
"""
|
||||
Generator that processes images in batches to optimize GPU-CPU transfer.
|
||||
|
||||
Args:
|
||||
image_sequence: A torch.Tensor or list of tensors/images
|
||||
batch_size: Number of frames to process at once for Tensor inputs
|
||||
|
||||
Yields:
|
||||
Numpy array for each batch or frame, contiguous and ready for ffmpeg
|
||||
"""
|
||||
# Optimized path for Tensor input
|
||||
if isinstance(image_sequence, torch.Tensor):
|
||||
total = len(image_sequence)
|
||||
for i in range(0, total, batch_size):
|
||||
# Process a chunk of frames on GPU/CPU together
|
||||
# This amortizes the overhead of kernel launches and synchronization
|
||||
batch = image_sequence[i:i+batch_size]
|
||||
batch_np = tensor_to_numpy_uint8(batch)
|
||||
# Yield the whole batch at once to optimize pipe writes
|
||||
yield np.ascontiguousarray(batch_np)
|
||||
else:
|
||||
# Fallback for list input (e.g. pingpong or mixed sources)
|
||||
# We process individually as stacking might be expensive if they are not already contiguous tensors
|
||||
for img in image_sequence:
|
||||
yield np.ascontiguousarray(tensor_to_numpy_uint8(img))
|
||||
@@ -0,0 +1,529 @@
|
||||
"""
|
||||
Video encoding utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides FFmpeg-based video encoding with fallback to PIL for GIF/WebP.
|
||||
"""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
from typing import List, Tuple, Optional, Iterator, Any, Callable
|
||||
from uuid import uuid4
|
||||
import numpy as np
|
||||
|
||||
|
||||
def detect_ffmpeg() -> Optional[str]:
|
||||
"""
|
||||
Detect FFmpeg binary location.
|
||||
|
||||
Returns:
|
||||
Path to FFmpeg executable, or None if not found
|
||||
"""
|
||||
ffmpeg_path = None
|
||||
|
||||
# Try imageio-ffmpeg first (common in Python environments)
|
||||
try:
|
||||
import imageio_ffmpeg
|
||||
ffmpeg_path = imageio_ffmpeg.get_ffmpeg_exe()
|
||||
print(f"Found ffmpeg via imageio_ffmpeg: {ffmpeg_path}")
|
||||
return ffmpeg_path
|
||||
except (ImportError, Exception):
|
||||
pass
|
||||
|
||||
# Try system PATH
|
||||
try:
|
||||
from shutil import which
|
||||
ffmpeg_path = which("ffmpeg")
|
||||
if ffmpeg_path:
|
||||
print(f"Found ffmpeg in system path: {ffmpeg_path}")
|
||||
return ffmpeg_path
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class FFmpegEncoder:
|
||||
"""
|
||||
FFmpeg-based video encoder supporting multiple formats.
|
||||
"""
|
||||
|
||||
def __init__(self, ffmpeg_path: Optional[str] = None):
|
||||
"""
|
||||
Initialize the encoder.
|
||||
|
||||
Args:
|
||||
ffmpeg_path: Path to FFmpeg executable (auto-detected if None)
|
||||
"""
|
||||
self.ffmpeg_path = ffmpeg_path or detect_ffmpeg()
|
||||
if not self.ffmpeg_path:
|
||||
raise RuntimeError("FFmpeg not found. Install ffmpeg or imageio-ffmpeg.")
|
||||
|
||||
def encode(
|
||||
self,
|
||||
images: List[np.ndarray],
|
||||
output_path: str,
|
||||
format_ext: str,
|
||||
frame_rate: float = 24.0,
|
||||
quality: int = 85,
|
||||
lossless: bool = False,
|
||||
loop_count: int = 0,
|
||||
codec: Optional[str] = None,
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Encode images to video using FFmpeg.
|
||||
|
||||
Args:
|
||||
images: List of numpy arrays (H, W, C) in uint8 format
|
||||
output_path: Output file path
|
||||
format_ext: Output format extension (mp4, webm, gif)
|
||||
frame_rate: Frame rate in FPS
|
||||
quality: Quality level 1-100
|
||||
lossless: Use lossless encoding if supported
|
||||
loop_count: Number of loops (0 = infinite for GIF)
|
||||
codec: Specific codec to use (h264, h265, vp9, etc.)
|
||||
progress_callback: Optional callback(current, total) for progress
|
||||
|
||||
Returns:
|
||||
Path to the encoded file
|
||||
"""
|
||||
if not images:
|
||||
raise ValueError("No images provided for encoding")
|
||||
|
||||
# Get dimensions from first image
|
||||
height, width = images[0].shape[:2]
|
||||
has_alpha = images[0].shape[2] == 4 if len(images[0].shape) > 2 else False
|
||||
|
||||
# Determine input pixel format
|
||||
i_pix_fmt = "rgba" if has_alpha else "rgb24"
|
||||
dimensions = f"{width}x{height}"
|
||||
|
||||
# Build FFmpeg arguments
|
||||
args = self._build_ffmpeg_args(
|
||||
format_ext=format_ext,
|
||||
dimensions=dimensions,
|
||||
frame_rate=frame_rate,
|
||||
quality=quality,
|
||||
lossless=lossless,
|
||||
loop_count=loop_count,
|
||||
i_pix_fmt=i_pix_fmt,
|
||||
has_alpha=has_alpha,
|
||||
codec=codec,
|
||||
output_path=output_path
|
||||
)
|
||||
|
||||
# Execute encoding
|
||||
self._execute_encoding(args, images, i_pix_fmt, progress_callback)
|
||||
|
||||
return output_path
|
||||
|
||||
def _build_ffmpeg_args(
|
||||
self,
|
||||
format_ext: str,
|
||||
dimensions: str,
|
||||
frame_rate: float,
|
||||
quality: int,
|
||||
lossless: bool,
|
||||
loop_count: int,
|
||||
i_pix_fmt: str,
|
||||
has_alpha: bool,
|
||||
codec: Optional[str],
|
||||
output_path: str
|
||||
) -> List[str]:
|
||||
"""Build FFmpeg command arguments."""
|
||||
|
||||
# Loop arguments
|
||||
loop_args = []
|
||||
if format_ext == "gif":
|
||||
loop_args = ["-loop", "0" if loop_count == 0 else str(loop_count)]
|
||||
|
||||
# Base input arguments
|
||||
args = [
|
||||
self.ffmpeg_path, "-v", "error",
|
||||
"-f", "rawvideo",
|
||||
"-pix_fmt", i_pix_fmt,
|
||||
"-s", dimensions,
|
||||
"-r", str(frame_rate),
|
||||
"-i", "-"
|
||||
] + loop_args
|
||||
|
||||
# Format-specific encoding arguments
|
||||
if format_ext == "gif":
|
||||
args.extend(self._get_gif_args(quality))
|
||||
elif format_ext == "mp4":
|
||||
args.extend(self._get_mp4_args(quality, lossless, codec))
|
||||
elif format_ext == "webm":
|
||||
args.extend(self._get_webm_args(quality, lossless, has_alpha))
|
||||
else:
|
||||
# Default to MP4-like encoding
|
||||
args.extend(self._get_mp4_args(quality, lossless, codec))
|
||||
|
||||
# Add output path
|
||||
args.extend(["-y", output_path])
|
||||
|
||||
return args
|
||||
|
||||
def _get_gif_args(self, quality: int) -> List[str]:
|
||||
"""Get FFmpeg arguments for GIF encoding."""
|
||||
# Use palettegen for better quality
|
||||
if quality >= 80:
|
||||
return [
|
||||
"-vf", "split[s0][s1];[s0]palettegen=max_colors=256:stats_mode=diff[p];[s1][p]paletteuse=dither=sierra2",
|
||||
"-f", "gif"
|
||||
]
|
||||
else:
|
||||
return [
|
||||
"-vf", "split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse",
|
||||
"-f", "gif"
|
||||
]
|
||||
|
||||
def _get_mp4_args(self, quality: int, lossless: bool, codec: Optional[str]) -> List[str]:
|
||||
"""Get FFmpeg arguments for MP4 encoding."""
|
||||
args = []
|
||||
|
||||
# Determine codec
|
||||
use_h265 = codec == "h265" or codec == "hevc"
|
||||
|
||||
if lossless:
|
||||
if use_h265:
|
||||
args.extend(["-c:v", "libx265", "-x265-params", "lossless=1"])
|
||||
else:
|
||||
args.extend(["-c:v", "libx264", "-crf", "0"])
|
||||
else:
|
||||
# Map quality (1-100) to CRF (51-0 for h264, lower is better)
|
||||
crf = int(51 - (quality / 100 * 33)) # Maps 1->51, 100->18
|
||||
|
||||
if use_h265:
|
||||
args.extend(["-c:v", "libx265", "-crf", str(crf + 5)]) # H.265 uses different CRF scale
|
||||
else:
|
||||
args.extend(["-c:v", "libx264", "-crf", str(crf)])
|
||||
|
||||
# Always use yuv420p for Discord compatibility
|
||||
args.extend(["-pix_fmt", "yuv420p", "-movflags", "faststart"])
|
||||
|
||||
return args
|
||||
|
||||
def _get_webm_args(self, quality: int, lossless: bool, has_alpha: bool) -> List[str]:
|
||||
"""Get FFmpeg arguments for WebM encoding."""
|
||||
args = ["-c:v", "libvpx-vp9"]
|
||||
|
||||
if lossless:
|
||||
args.extend(["-lossless", "1"])
|
||||
else:
|
||||
# Map quality to CRF (63-0 for VP9)
|
||||
crf = int(63 - (quality / 100 * 33)) # Maps 1->63, 100->30
|
||||
args.extend(["-crf", str(crf), "-b:v", "0"])
|
||||
|
||||
# Pixel format - support alpha if present
|
||||
pix_fmt = "yuva420p" if has_alpha else "yuv420p"
|
||||
args.extend(["-pix_fmt", pix_fmt])
|
||||
|
||||
# VP9 threading
|
||||
args.extend(["-row-mt", "1"])
|
||||
|
||||
return args
|
||||
|
||||
def _execute_encoding(
|
||||
self,
|
||||
args: List[str],
|
||||
images: List[np.ndarray],
|
||||
i_pix_fmt: str,
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None
|
||||
) -> None:
|
||||
"""Execute FFmpeg process and feed frames."""
|
||||
total_frames = len(images)
|
||||
|
||||
# Create image chunk iterator for memory efficiency
|
||||
def image_chunks() -> Iterator[bytes]:
|
||||
for i, img in enumerate(images):
|
||||
# Ensure contiguous array for subprocess
|
||||
chunk = np.ascontiguousarray(img)
|
||||
if progress_callback:
|
||||
progress_callback(i + 1, total_frames)
|
||||
yield chunk.tobytes()
|
||||
|
||||
# Start FFmpeg process
|
||||
process = subprocess.Popen(
|
||||
args,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE
|
||||
)
|
||||
|
||||
# Feed frames
|
||||
try:
|
||||
for chunk in image_chunks():
|
||||
process.stdin.write(chunk)
|
||||
process.stdin.close()
|
||||
process.wait()
|
||||
|
||||
if process.returncode != 0:
|
||||
stderr = process.stderr.read().decode('utf-8', errors='ignore')
|
||||
raise RuntimeError(f"FFmpeg encoding failed: {stderr}")
|
||||
finally:
|
||||
if process.stdin:
|
||||
process.stdin.close()
|
||||
if process.stdout:
|
||||
process.stdout.close()
|
||||
if process.stderr:
|
||||
process.stderr.close()
|
||||
|
||||
|
||||
class PILEncoder:
|
||||
"""
|
||||
PIL-based encoder for GIF and WebP formats.
|
||||
Fallback when FFmpeg is not available.
|
||||
"""
|
||||
|
||||
def encode(
|
||||
self,
|
||||
images: List[Any], # PIL Images or numpy arrays
|
||||
output_path: str,
|
||||
format_ext: str,
|
||||
frame_rate: float = 24.0,
|
||||
quality: int = 85,
|
||||
lossless: bool = False,
|
||||
loop_count: int = 0,
|
||||
tensor_to_numpy_func: Optional[Callable] = None
|
||||
) -> str:
|
||||
"""
|
||||
Encode images using PIL.
|
||||
|
||||
Args:
|
||||
images: List of PIL Images or numpy arrays
|
||||
output_path: Output file path
|
||||
format_ext: Output format (gif, webp)
|
||||
frame_rate: Frame rate in FPS
|
||||
quality: Quality level 1-100
|
||||
lossless: Use lossless encoding for WebP
|
||||
loop_count: Number of loops (0 = infinite)
|
||||
tensor_to_numpy_func: Optional function to convert tensors to numpy
|
||||
|
||||
Returns:
|
||||
Path to the encoded file
|
||||
"""
|
||||
from PIL import Image
|
||||
|
||||
# Convert to PIL images if needed
|
||||
pil_images = []
|
||||
for img in images:
|
||||
if hasattr(img, 'shape'): # numpy array or tensor
|
||||
if tensor_to_numpy_func and hasattr(img, 'cpu'):
|
||||
img = tensor_to_numpy_func(img)
|
||||
elif hasattr(img, 'numpy'):
|
||||
img = img.numpy()
|
||||
pil_images.append(Image.fromarray(img.astype(np.uint8)))
|
||||
else:
|
||||
pil_images.append(img)
|
||||
|
||||
if not pil_images:
|
||||
raise ValueError("No images provided for encoding")
|
||||
|
||||
# Calculate frame duration in milliseconds
|
||||
duration = int(1000 / frame_rate)
|
||||
|
||||
if format_ext.lower() == "gif":
|
||||
self._encode_gif(pil_images, output_path, duration, loop_count)
|
||||
elif format_ext.lower() == "webp":
|
||||
self._encode_webp(pil_images, output_path, duration, loop_count, quality, lossless)
|
||||
else:
|
||||
# Single frame fallback
|
||||
pil_images[0].save(output_path, format=format_ext.upper())
|
||||
|
||||
return output_path
|
||||
|
||||
def _encode_gif(
|
||||
self,
|
||||
images: List[Any],
|
||||
output_path: str,
|
||||
duration: int,
|
||||
loop_count: int
|
||||
) -> None:
|
||||
"""Encode images as GIF."""
|
||||
durations = [duration] * len(images)
|
||||
images[0].save(
|
||||
output_path,
|
||||
format="GIF",
|
||||
append_images=images[1:] if len(images) > 1 else [],
|
||||
save_all=True,
|
||||
duration=durations,
|
||||
loop=0 if loop_count == 0 else loop_count,
|
||||
optimize=False
|
||||
)
|
||||
|
||||
def _encode_webp(
|
||||
self,
|
||||
images: List[Any],
|
||||
output_path: str,
|
||||
duration: int,
|
||||
loop_count: int,
|
||||
quality: int,
|
||||
lossless: bool
|
||||
) -> None:
|
||||
"""Encode images as WebP."""
|
||||
save_kwargs = {
|
||||
"format": "WEBP",
|
||||
"append_images": images[1:] if len(images) > 1 else [],
|
||||
"save_all": True,
|
||||
"duration": duration,
|
||||
"loop": 0 if loop_count == 0 else loop_count,
|
||||
}
|
||||
|
||||
if lossless:
|
||||
save_kwargs["lossless"] = True
|
||||
else:
|
||||
save_kwargs["quality"] = quality
|
||||
|
||||
images[0].save(output_path, **save_kwargs)
|
||||
|
||||
|
||||
def optimize_video_for_discord(
|
||||
input_file: str,
|
||||
ffmpeg_path: str,
|
||||
temp_dir: str
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Create a Discord-optimized version of a video file.
|
||||
|
||||
Args:
|
||||
input_file: Path to the input video file
|
||||
ffmpeg_path: Path to FFmpeg executable
|
||||
temp_dir: Directory for temporary files
|
||||
|
||||
Returns:
|
||||
Path to the optimized file, or None if optimization failed
|
||||
"""
|
||||
format_ext = os.path.splitext(input_file)[1].lstrip('.').lower()
|
||||
|
||||
# Use mkstemp for secure temporary file creation with restricted permissions (0600)
|
||||
# This prevents race conditions and ensures other users can't read the temp file
|
||||
fd, discord_optimized_file = tempfile.mkstemp(
|
||||
suffix=f".{format_ext}",
|
||||
prefix="discord_optimized_",
|
||||
dir=temp_dir
|
||||
)
|
||||
os.close(fd) # Close file descriptor immediately so FFmpeg can write to it
|
||||
|
||||
success = False
|
||||
try:
|
||||
if format_ext == "mp4":
|
||||
optimize_args = [
|
||||
ffmpeg_path, "-i", input_file,
|
||||
"-c:v", "libx264", "-pix_fmt", "yuv420p",
|
||||
"-movflags", "faststart", "-preset", "fast",
|
||||
"-profile:v", "baseline", "-level", "3.0",
|
||||
"-crf", "23",
|
||||
"-c:a", "aac", "-b:a", "128k",
|
||||
"-y", discord_optimized_file
|
||||
]
|
||||
elif format_ext == "webm":
|
||||
optimize_args = [
|
||||
ffmpeg_path, "-i", input_file,
|
||||
"-c:v", "libvpx-vp9",
|
||||
"-pix_fmt", "yuv420p",
|
||||
"-crf", "30", "-b:v", "0",
|
||||
"-deadline", "good",
|
||||
"-c:a", "libopus", "-b:a", "96k",
|
||||
"-y", discord_optimized_file
|
||||
]
|
||||
elif format_ext == "gif":
|
||||
optimize_args = [
|
||||
ffmpeg_path, "-i", input_file,
|
||||
"-vf", "fps=15,scale=trunc(iw/2)*2:trunc(ih/2)*2",
|
||||
"-y", discord_optimized_file
|
||||
]
|
||||
else:
|
||||
print(f"No optimization rules for format: {format_ext}")
|
||||
return None
|
||||
|
||||
print(f"Creating Discord-optimized version of {format_ext.upper()} file...")
|
||||
result = subprocess.run(
|
||||
optimize_args,
|
||||
capture_output=True,
|
||||
text=True
|
||||
)
|
||||
|
||||
if result.returncode == 0 and os.path.exists(discord_optimized_file):
|
||||
print(f"Discord-optimized file created: {discord_optimized_file}")
|
||||
success = True
|
||||
return discord_optimized_file
|
||||
else:
|
||||
print(f"Optimization failed: {result.stderr}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error during Discord optimization: {e}")
|
||||
return None
|
||||
|
||||
finally:
|
||||
# Clean up temp file if optimization failed or wasn't supported
|
||||
if not success and os.path.exists(discord_optimized_file):
|
||||
try:
|
||||
os.remove(discord_optimized_file)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def mux_audio_to_video(
|
||||
video_path: str,
|
||||
audio_waveform: np.ndarray,
|
||||
sample_rate: int,
|
||||
format_ext: str,
|
||||
ffmpeg_path: str,
|
||||
output_path: str,
|
||||
channels: int = 2
|
||||
) -> bool:
|
||||
"""
|
||||
Mux audio into a video file.
|
||||
|
||||
Args:
|
||||
video_path: Path to the video file
|
||||
audio_waveform: Audio data as numpy array
|
||||
sample_rate: Audio sample rate
|
||||
format_ext: Video format extension
|
||||
ffmpeg_path: Path to FFmpeg executable
|
||||
output_path: Output path for the muxed file
|
||||
channels: Number of audio channels
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Determine audio codec based on format
|
||||
if format_ext == "mp4":
|
||||
audio_pass = ["-c:a", "aac", "-b:a", "192k"]
|
||||
elif format_ext == "webm":
|
||||
audio_pass = ["-c:a", "libopus", "-b:a", "128k"]
|
||||
else:
|
||||
audio_pass = ["-c:a", "libopus", "-b:a", "128k"]
|
||||
|
||||
mux_args = [
|
||||
ffmpeg_path, "-v", "error", "-y",
|
||||
"-i", video_path,
|
||||
"-ar", str(sample_rate),
|
||||
"-ac", str(channels),
|
||||
"-f", "f32le",
|
||||
"-i", "-",
|
||||
"-c:v", "copy"
|
||||
] + audio_pass + ["-shortest", output_path]
|
||||
|
||||
# Ensure contiguous array for subprocess
|
||||
audio_data = np.ascontiguousarray(audio_waveform)
|
||||
|
||||
result = subprocess.run(
|
||||
mux_args,
|
||||
input=memoryview(audio_data),
|
||||
capture_output=True
|
||||
)
|
||||
|
||||
if result.returncode == 0:
|
||||
print(f"Successfully muxed audio to video: {output_path}")
|
||||
return True
|
||||
else:
|
||||
print(f"Audio muxing failed: {result.stderr.decode('utf-8', errors='ignore')}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error muxing audio: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,147 @@
|
||||
"""
|
||||
Path utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides functions for handling output directories and file paths.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_output_directory(
|
||||
save_output: bool,
|
||||
comfy_output_dir: str,
|
||||
temp_dir: str,
|
||||
subfolder: str = "discord_output"
|
||||
) -> str:
|
||||
"""
|
||||
Determine the appropriate output directory based on save settings.
|
||||
|
||||
Args:
|
||||
save_output: Whether files should be saved permanently
|
||||
comfy_output_dir: ComfyUI's output directory path
|
||||
temp_dir: ComfyUI's temporary directory path
|
||||
subfolder: Subfolder name within output directory (default: "discord_output")
|
||||
|
||||
Returns:
|
||||
Path to the destination directory
|
||||
"""
|
||||
if save_output:
|
||||
# Create output subfolder in the ComfyUI output directory
|
||||
dest_folder = os.path.join(comfy_output_dir, subfolder)
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
else:
|
||||
# Use ComfyUI's temporary directory for preview-only files
|
||||
dest_folder = temp_dir
|
||||
os.makedirs(dest_folder, exist_ok=True)
|
||||
print(f"Using temporary directory for preview: {dest_folder}")
|
||||
|
||||
return dest_folder
|
||||
|
||||
|
||||
def ensure_directory_exists(path: str) -> str:
|
||||
"""
|
||||
Ensure a directory exists, creating it if necessary.
|
||||
|
||||
Args:
|
||||
path: Directory path to ensure exists
|
||||
|
||||
Returns:
|
||||
The same path (for chaining)
|
||||
"""
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def get_unique_filepath(
|
||||
directory: str,
|
||||
filename: str,
|
||||
extension: str,
|
||||
counter: Optional[int] = None
|
||||
) -> str:
|
||||
"""
|
||||
Generate a unique filepath, optionally with a counter.
|
||||
|
||||
Args:
|
||||
directory: Base directory
|
||||
filename: Base filename (without extension)
|
||||
extension: File extension (with or without leading dot)
|
||||
counter: Optional counter to append to filename
|
||||
|
||||
Returns:
|
||||
Full filepath
|
||||
"""
|
||||
# Ensure extension has leading dot
|
||||
if not extension.startswith("."):
|
||||
extension = "." + extension
|
||||
|
||||
if counter is not None:
|
||||
full_filename = f"{filename}_{counter:05d}{extension}"
|
||||
else:
|
||||
full_filename = f"{filename}{extension}"
|
||||
|
||||
return os.path.join(directory, full_filename)
|
||||
|
||||
|
||||
def validate_path_is_safe(path: str, base_dir: Optional[str] = None) -> None:
|
||||
"""
|
||||
Validate that a path is safe to write to.
|
||||
|
||||
Checks:
|
||||
- If base_dir is provided, path is contained within base_dir (to prevent ../ traversal)
|
||||
- Path is not a symlink (to prevent overwriting targets)
|
||||
- Parent directories are not symlinks (to prevent path traversal via symlinks)
|
||||
|
||||
Args:
|
||||
path: File path to validate
|
||||
base_dir: Optional base directory to restrict path to
|
||||
|
||||
Raises:
|
||||
ValueError: If path is unsafe
|
||||
"""
|
||||
# Check if path is within base_dir
|
||||
if base_dir:
|
||||
abs_base = os.path.abspath(base_dir)
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
# Use commonpath to ensure path is within base_dir
|
||||
# We need to handle potential different drives on Windows which raises ValueError
|
||||
try:
|
||||
common = os.path.commonpath([abs_base, abs_path])
|
||||
except ValueError:
|
||||
# Raised if paths are on different drives
|
||||
raise ValueError(f"Security error: Path '{path}' is on a different drive than allowed directory '{base_dir}'.")
|
||||
|
||||
if common != abs_base:
|
||||
raise ValueError(f"Security error: Path '{path}' is outside the allowed directory '{base_dir}'.")
|
||||
|
||||
# Check if path exists and is a symlink
|
||||
if os.path.islink(path):
|
||||
raise ValueError(f"Security error: Output path '{path}' is a symlink. Overwriting symlinks is not allowed.")
|
||||
|
||||
# Verify parent directories
|
||||
# Walk up the tree to find the first existing directory
|
||||
current_dir = os.path.dirname(os.path.abspath(path))
|
||||
|
||||
# Safety valve to prevent infinite loops (though OS paths are finite)
|
||||
# We check existence. If it doesn't exist, we check if it's a broken symlink (islink returns True even for broken links)
|
||||
# Then move to parent.
|
||||
|
||||
while current_dir and current_dir != os.path.dirname(current_dir): # Until root
|
||||
if os.path.islink(current_dir):
|
||||
raise ValueError(f"Security error: Path component '{current_dir}' is a symlink. Writing through directory symlinks is not allowed.")
|
||||
|
||||
if os.path.exists(current_dir):
|
||||
# Once we hit an existing directory, we verify it matches its realpath
|
||||
# This catches hidden symlinks further up that might have been resolved by abspath but diverge in realpath
|
||||
real_dir = os.path.realpath(current_dir)
|
||||
abs_dir = os.path.abspath(current_dir)
|
||||
|
||||
if real_dir != abs_dir:
|
||||
raise ValueError(f"Security error: Path resolution mismatch for '{current_dir}'. "
|
||||
f"Symlinks in output paths are not allowed (Real: {real_dir}, Abs: {abs_dir}).")
|
||||
# If the existing ancestor is safe, we assume children created under it will be normal directories
|
||||
# (unless we have a race condition, but we can't solve that fully without openat)
|
||||
break
|
||||
|
||||
current_dir = os.path.dirname(current_dir)
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Workflow Manipulation Utilities
|
||||
|
||||
Provides sanitization, prompt extraction, and workflow building tools.
|
||||
"""
|
||||
|
||||
from .sanitizer import sanitize_json_for_export
|
||||
from .prompt_extractor import extract_prompts_from_workflow
|
||||
from .workflow_builder import WorkflowBuilder
|
||||
|
||||
__all__ = [
|
||||
'sanitize_json_for_export',
|
||||
'extract_prompts_from_workflow',
|
||||
'WorkflowBuilder',
|
||||
]
|
||||
@@ -0,0 +1,270 @@
|
||||
"""
|
||||
Prompt Extraction Utilities for ComfyUI-DiscordSend
|
||||
|
||||
Extracts positive and negative prompts from ComfyUI workflow data.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
|
||||
# Common negative prompt indicators
|
||||
NEGATIVE_INDICATORS = [
|
||||
"bad quality", "deformed", "blurry", "low quality", "worst quality",
|
||||
"ugly", "disfigured", "low res", "poorly drawn", "mutation",
|
||||
"extra limbs", "bad anatomy", "watermark", "text", "signature"
|
||||
]
|
||||
|
||||
# Node types that can contain prompts
|
||||
PROMPT_NODE_TYPES = [
|
||||
"CLIPTextEncode", # Standard SD 1.5 prompt node
|
||||
"SDXLPromptEncoder", # SDXL prompt encoder
|
||||
"SDXLTextEncode", # Another SDXL text node
|
||||
]
|
||||
|
||||
|
||||
def extract_prompts_from_workflow(workflow_data: Any) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
Extract positive and negative prompts from workflow data.
|
||||
|
||||
Analyzes ComfyUI workflow structure to find prompt nodes (CLIPTextEncode,
|
||||
SDXLPromptEncoder, SDXLTextEncode, and other text encoding nodes) and
|
||||
determine which contains the positive vs negative prompt.
|
||||
|
||||
Args:
|
||||
workflow_data: The workflow data dictionary or JSON string
|
||||
|
||||
Returns:
|
||||
A tuple of (positive_prompt, negative_prompt) or (None, None) if not found
|
||||
"""
|
||||
if workflow_data is None:
|
||||
return None, None
|
||||
|
||||
# Convert string to dict if necessary
|
||||
if isinstance(workflow_data, str):
|
||||
try:
|
||||
data = json.loads(workflow_data)
|
||||
except json.JSONDecodeError:
|
||||
return None, None
|
||||
else:
|
||||
data = workflow_data
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return None, None
|
||||
|
||||
positive_prompt = None
|
||||
negative_prompt = None
|
||||
|
||||
# Find prompt nodes (CLIPTextEncode, SDXL nodes, etc.)
|
||||
if "nodes" in data:
|
||||
nodes = data["nodes"]
|
||||
else:
|
||||
# Check if it's API format (dict of nodes)
|
||||
nodes = data
|
||||
|
||||
prompt_nodes = _find_prompt_nodes(nodes)
|
||||
|
||||
if not prompt_nodes:
|
||||
return None, None
|
||||
|
||||
# Determine positive/negative based on content and structure
|
||||
if len(prompt_nodes) == 1:
|
||||
# Single prompt node - assume it's the positive prompt
|
||||
positive_prompt = _get_prompt_text(prompt_nodes[0])
|
||||
elif len(prompt_nodes) >= 2:
|
||||
# Multiple prompt nodes - need to determine which is which
|
||||
positive_prompt, negative_prompt = _classify_prompts(prompt_nodes, data)
|
||||
|
||||
# Ensure we return empty string for negative if we have positive but not negative
|
||||
if positive_prompt is not None and negative_prompt is None:
|
||||
negative_prompt = ""
|
||||
|
||||
return positive_prompt, negative_prompt
|
||||
|
||||
|
||||
def _find_prompt_nodes(nodes: Union[List, Dict]) -> List[Dict]:
|
||||
"""Find all prompt nodes (CLIPTextEncode, SDXL nodes, etc.) in the workflow."""
|
||||
prompt_nodes = []
|
||||
|
||||
if isinstance(nodes, list):
|
||||
for node in nodes:
|
||||
if _is_prompt_node(node):
|
||||
prompt_nodes.append(node)
|
||||
elif isinstance(nodes, dict):
|
||||
for node_id, node in nodes.items():
|
||||
if _is_prompt_node(node):
|
||||
node_copy = dict(node)
|
||||
node_copy["id"] = node_id
|
||||
prompt_nodes.append(node_copy)
|
||||
|
||||
return prompt_nodes
|
||||
|
||||
|
||||
def _is_prompt_node(node: Any) -> bool:
|
||||
"""Check if a node is a prompt node (CLIPTextEncode, SDXL, etc.) with valid text."""
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
|
||||
# Handle both Workflow format (type) and API format (class_type)
|
||||
node_type = node.get("type") or node.get("class_type")
|
||||
|
||||
# Check against known prompt node types
|
||||
if node_type not in PROMPT_NODE_TYPES:
|
||||
# Also check for dynamic text/prompt nodes (e.g., custom nodes)
|
||||
if node_type and ("Text" in node_type and ("Encode" in node_type or "Prompt" in node_type)):
|
||||
pass # Allow these through
|
||||
else:
|
||||
return False
|
||||
|
||||
# Check for text in either widgets_values (Workflow) or inputs (API)
|
||||
text = _get_prompt_text(node)
|
||||
return text is not None
|
||||
|
||||
|
||||
def _get_prompt_text(node: Dict) -> Optional[str]:
|
||||
"""Extract the prompt text from a prompt node."""
|
||||
# Workflow format (widgets_values)
|
||||
widgets = node.get("widgets_values", [])
|
||||
if isinstance(widgets, list) and len(widgets) > 0 and isinstance(widgets[0], str):
|
||||
return widgets[0]
|
||||
|
||||
# API format (inputs.text)
|
||||
inputs = node.get("inputs", {})
|
||||
if isinstance(inputs, dict) and isinstance(inputs.get("text"), str):
|
||||
return inputs.get("text")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _classify_prompts(prompt_nodes: List[Dict], workflow_data: Dict) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
Classify which prompt nodes contain positive vs negative prompts.
|
||||
|
||||
Uses multiple heuristics:
|
||||
1. Content analysis (negative prompts often contain quality-related terms)
|
||||
2. Connection analysis (traces connections to sampler nodes)
|
||||
"""
|
||||
if not prompt_nodes:
|
||||
return None, None
|
||||
|
||||
# First pass: Score all nodes based on content
|
||||
node_scores = []
|
||||
for node in prompt_nodes:
|
||||
prompt_text = _get_prompt_text(node)
|
||||
# Skip empty or None text
|
||||
if not prompt_text or not prompt_text.strip():
|
||||
continue
|
||||
|
||||
prompt_lower = prompt_text.lower()
|
||||
score = sum(1 for indicator in NEGATIVE_INDICATORS if indicator in prompt_lower)
|
||||
node_scores.append((node, score, prompt_text))
|
||||
|
||||
if not node_scores:
|
||||
return None, None
|
||||
|
||||
# Sort by score descending (highest score first for negative)
|
||||
node_scores.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
positive_prompt = None
|
||||
negative_prompt = None
|
||||
|
||||
# Highest score is likely negative (if score > 0)
|
||||
if node_scores[0][1] > 0:
|
||||
negative_prompt = node_scores[0][2]
|
||||
# The node with the lowest score is likely the positive prompt
|
||||
# Since it's sorted descending, pick the last one
|
||||
positive_prompt = node_scores[-1][2]
|
||||
|
||||
# If we have more than 2 nodes and the last one is NOT the only zero-score node,
|
||||
# we might need more logic, but for 2 nodes, this is correct.
|
||||
# If positive and negative are the same (only 1 node with score > 0), clear positive.
|
||||
if positive_prompt == negative_prompt:
|
||||
positive_prompt = None
|
||||
else:
|
||||
# All scores are 0, use connection analysis
|
||||
positive_prompt, negative_prompt = _classify_by_connections(
|
||||
prompt_nodes, workflow_data, None, None
|
||||
)
|
||||
|
||||
# Fallback: if we still can't determine, use first two nodes
|
||||
if positive_prompt is None and negative_prompt is None and len(prompt_nodes) >= 2:
|
||||
# Convention: assume first is positive, second is negative
|
||||
positive_prompt = _get_prompt_text(prompt_nodes[0])
|
||||
negative_prompt = _get_prompt_text(prompt_nodes[1])
|
||||
elif positive_prompt is None and negative_prompt is not None:
|
||||
# Find the other prompt
|
||||
for _, _, text in node_scores:
|
||||
if text != negative_prompt:
|
||||
positive_prompt = text
|
||||
break
|
||||
elif negative_prompt is None and positive_prompt is not None:
|
||||
# Find the other prompt
|
||||
for _, _, text in node_scores:
|
||||
if text != positive_prompt:
|
||||
negative_prompt = text
|
||||
break
|
||||
|
||||
return positive_prompt, negative_prompt
|
||||
|
||||
|
||||
def _classify_by_connections(
|
||||
prompt_nodes: List[Dict],
|
||||
workflow_data: Dict,
|
||||
existing_positive: Optional[str],
|
||||
existing_negative: Optional[str]
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""Try to classify prompts by analyzing node connections to samplers."""
|
||||
positive = existing_positive
|
||||
negative = existing_negative
|
||||
|
||||
links = workflow_data.get("links", [])
|
||||
if not isinstance(links, list):
|
||||
return positive, negative
|
||||
|
||||
nodes = workflow_data.get("nodes", [])
|
||||
|
||||
# Find sampler nodes
|
||||
samplers = []
|
||||
if isinstance(nodes, list):
|
||||
for node in nodes:
|
||||
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
|
||||
samplers.append(node)
|
||||
elif isinstance(nodes, dict):
|
||||
for node_id, node in nodes.items():
|
||||
if isinstance(node, dict) and "KSampler" in node.get("type", ""):
|
||||
node_copy = dict(node)
|
||||
node_copy["id"] = node_id
|
||||
samplers.append(node_copy)
|
||||
|
||||
if not samplers:
|
||||
return positive, negative
|
||||
|
||||
# Trace connections from CLIP nodes to samplers
|
||||
for link in links:
|
||||
if len(link) < 4:
|
||||
continue
|
||||
|
||||
from_node_id = link[0]
|
||||
to_node_id = link[2]
|
||||
to_slot = link[3]
|
||||
|
||||
# Find matching CLIP node and sampler
|
||||
for clip_node in prompt_nodes:
|
||||
clip_id = clip_node.get("id")
|
||||
if clip_id == from_node_id:
|
||||
for sampler in samplers:
|
||||
if sampler.get("id") == to_node_id:
|
||||
# Determine slot type
|
||||
inputs = sampler.get("inputs", {})
|
||||
if isinstance(inputs, dict):
|
||||
input_keys = list(inputs.keys())
|
||||
if to_slot < len(input_keys):
|
||||
input_name = input_keys[to_slot].lower()
|
||||
prompt_text = _get_prompt_text(clip_node)
|
||||
|
||||
if "positive" in input_name and positive is None:
|
||||
positive = prompt_text
|
||||
elif "negative" in input_name and negative is None:
|
||||
negative = prompt_text
|
||||
|
||||
return positive, negative
|
||||
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
JSON Sanitization Utilities for ComfyUI-DiscordSend
|
||||
|
||||
Provides comprehensive sanitization of sensitive data (webhook URLs, GitHub tokens)
|
||||
from JSON data structures before export or sharing.
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
|
||||
# Patterns for detecting sensitive data
|
||||
# Pre-compile regex for faster matching
|
||||
WEBHOOK_REGEX = re.compile(
|
||||
r"(discord\.com/api/webhooks|discordapp\.com/api/webhooks)",
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
GITHUB_TOKEN_PREFIXES = (
|
||||
"ghp_", # GitHub personal access token
|
||||
"github_pat_", # GitHub personal access token (new format)
|
||||
"gho_", # GitHub OAuth token
|
||||
"ghs_", # GitHub service token
|
||||
"ghu_", # GitHub user-to-server token
|
||||
)
|
||||
|
||||
|
||||
def is_webhook_url(value: str) -> bool:
|
||||
"""Check if a string appears to be a Discord webhook URL."""
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
|
||||
if WEBHOOK_REGEX.search(value):
|
||||
return True
|
||||
|
||||
# Also check for generic webhook patterns in URLs
|
||||
if value.startswith("http") and ("webhook" in value.lower() or "discord" in value.lower()):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def is_github_token(value: str) -> bool:
|
||||
"""Check if a string appears to be a GitHub token."""
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
|
||||
return value.startswith(GITHUB_TOKEN_PREFIXES)
|
||||
|
||||
|
||||
def is_potential_token(value: str, context_type: str = "") -> bool:
|
||||
"""Check if a string might be a token based on length and context."""
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
|
||||
# Long alphanumeric strings in GitHub-related contexts
|
||||
if len(value) >= 40 and "github" in context_type.lower():
|
||||
if any(c.isalnum() for c in value):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def sanitize_string(value: str, context_type: str = "") -> str:
|
||||
"""
|
||||
Sanitize a single string value.
|
||||
|
||||
Args:
|
||||
value: The string to sanitize
|
||||
context_type: Optional context (e.g., node type) to help identify tokens
|
||||
|
||||
Returns:
|
||||
Empty string if sensitive data detected, original value otherwise
|
||||
"""
|
||||
if is_webhook_url(value):
|
||||
return ""
|
||||
|
||||
if is_github_token(value):
|
||||
return ""
|
||||
|
||||
if is_potential_token(value, context_type):
|
||||
return ""
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def sanitize_widget_values(widgets: List, node_type: str = "") -> List:
|
||||
"""
|
||||
Sanitize a list of widget values from a ComfyUI node.
|
||||
|
||||
Args:
|
||||
widgets: List of widget values
|
||||
node_type: The type of node for context
|
||||
|
||||
Returns:
|
||||
Sanitized list with sensitive values replaced with empty strings
|
||||
"""
|
||||
result = []
|
||||
for value in widgets:
|
||||
if isinstance(value, str):
|
||||
result.append(sanitize_string(value, node_type))
|
||||
elif isinstance(value, dict):
|
||||
result.append(sanitize_dict(value))
|
||||
elif isinstance(value, list):
|
||||
result.append(sanitize_widget_values(value, node_type))
|
||||
else:
|
||||
result.append(value)
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_node_inputs(inputs: Dict, node_type: str = "") -> Dict:
|
||||
"""
|
||||
Sanitize the inputs dictionary of a ComfyUI node.
|
||||
|
||||
Args:
|
||||
inputs: Dictionary of node inputs
|
||||
node_type: The type of node for context
|
||||
|
||||
Returns:
|
||||
Sanitized dictionary
|
||||
"""
|
||||
result = {}
|
||||
for key, value in inputs.items():
|
||||
if key in ("webhook_url", "github_token"):
|
||||
result[key] = ""
|
||||
elif isinstance(value, str):
|
||||
result[key] = sanitize_string(value, node_type)
|
||||
elif isinstance(value, dict):
|
||||
result[key] = sanitize_node_inputs(value, node_type)
|
||||
elif isinstance(value, list):
|
||||
result[key] = sanitize_widget_values(value, node_type)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_node(node: Any) -> Any:
|
||||
"""
|
||||
Sanitize a single ComfyUI node.
|
||||
|
||||
Args:
|
||||
node: The node dictionary or value
|
||||
|
||||
Returns:
|
||||
Sanitized node dictionary or value
|
||||
"""
|
||||
if not isinstance(node, dict):
|
||||
if isinstance(node, str):
|
||||
return sanitize_string(node)
|
||||
return node
|
||||
|
||||
result = {}
|
||||
node_type = node.get("type", "")
|
||||
|
||||
for key, value in node.items():
|
||||
# Handle known sensitive keys
|
||||
if key in ("webhook_url", "github_token"):
|
||||
result[key] = ""
|
||||
continue
|
||||
|
||||
# Context-aware sanitization for inputs and widgets
|
||||
if key == "inputs" and isinstance(value, dict):
|
||||
result[key] = sanitize_node_inputs(value, node_type)
|
||||
elif key == "widgets_values" and isinstance(value, list):
|
||||
result[key] = sanitize_widget_values(value, node_type)
|
||||
else:
|
||||
# Generic sanitization for other fields
|
||||
if isinstance(value, dict):
|
||||
result[key] = sanitize_dict(value)
|
||||
elif isinstance(value, list):
|
||||
result[key] = sanitize_list(value)
|
||||
elif isinstance(value, str):
|
||||
result[key] = sanitize_string(value)
|
||||
else:
|
||||
result[key] = value
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_dict(data: Dict) -> Dict:
|
||||
"""
|
||||
Recursively sanitize a dictionary.
|
||||
|
||||
Args:
|
||||
data: Dictionary to sanitize
|
||||
|
||||
Returns:
|
||||
Sanitized dictionary
|
||||
"""
|
||||
result = {}
|
||||
|
||||
# Check if this is a workflow object with nodes
|
||||
# We want to process nodes specifically using sanitize_node to ensure correct context
|
||||
# and avoid double-processing (once as generic dict/list, once as nodes)
|
||||
is_workflow = "nodes" in data and isinstance(data["nodes"], (list, dict))
|
||||
|
||||
for key, value in data.items():
|
||||
# Handle known sensitive keys
|
||||
if key in ("webhook_url", "github_token"):
|
||||
result[key] = ""
|
||||
continue
|
||||
|
||||
# Special handling for "nodes" in workflow
|
||||
if is_workflow and key == "nodes":
|
||||
if isinstance(value, list):
|
||||
result[key] = [sanitize_node(n) for n in value]
|
||||
elif isinstance(value, dict):
|
||||
result[key] = {k: sanitize_node(v) for k, v in value.items()}
|
||||
else:
|
||||
# Fallback if nodes is neither list nor dict (unlikely)
|
||||
result[key] = value
|
||||
continue
|
||||
|
||||
# Handle nested structures
|
||||
if isinstance(value, dict):
|
||||
result[key] = sanitize_dict(value)
|
||||
elif isinstance(value, list):
|
||||
result[key] = sanitize_list(value)
|
||||
elif isinstance(value, str):
|
||||
result[key] = sanitize_string(value)
|
||||
else:
|
||||
result[key] = value
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_list(data: List) -> List:
|
||||
"""
|
||||
Recursively sanitize a list.
|
||||
|
||||
Args:
|
||||
data: List to sanitize
|
||||
|
||||
Returns:
|
||||
Sanitized list
|
||||
"""
|
||||
result = []
|
||||
for item in data:
|
||||
if isinstance(item, dict):
|
||||
result.append(sanitize_dict(item))
|
||||
elif isinstance(item, list):
|
||||
result.append(sanitize_list(item))
|
||||
elif isinstance(item, str):
|
||||
result.append(sanitize_string(item))
|
||||
else:
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
|
||||
def sanitize_json_for_export(json_data: Any) -> Any:
|
||||
"""
|
||||
Remove sensitive webhook data and GitHub tokens from JSON data to protect user security.
|
||||
|
||||
This is the main entry point for sanitization. It handles:
|
||||
- Discord webhook URLs (multiple formats)
|
||||
- GitHub personal access tokens (all variants)
|
||||
- Nested dictionaries and lists
|
||||
- ComfyUI workflow structures
|
||||
|
||||
Args:
|
||||
json_data: The JSON data object (dict, list, or string) to sanitize
|
||||
|
||||
Returns:
|
||||
The sanitized JSON data with sensitive information removed
|
||||
"""
|
||||
if json_data is None:
|
||||
return None
|
||||
|
||||
# Handle string input (may be JSON string or plain string)
|
||||
if isinstance(json_data, str):
|
||||
try:
|
||||
data = json.loads(json_data)
|
||||
result = sanitize_json_for_export(data)
|
||||
return json.dumps(result)
|
||||
except json.JSONDecodeError:
|
||||
# Not valid JSON - check if it's sensitive data directly
|
||||
return sanitize_string(json_data)
|
||||
|
||||
# Handle dictionary
|
||||
if isinstance(json_data, dict):
|
||||
return sanitize_dict(json_data)
|
||||
|
||||
# Handle list
|
||||
if isinstance(json_data, list):
|
||||
return sanitize_list(json_data)
|
||||
|
||||
# Return other types as-is
|
||||
return json_data
|
||||
@@ -0,0 +1,146 @@
|
||||
import json
|
||||
import random
|
||||
import logging
|
||||
from typing import Dict, Any, Optional, Tuple, List
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class WorkflowBuilder:
|
||||
"""Helper to manipulate ComfyUI workflow JSONs."""
|
||||
|
||||
def __init__(self, workflow_json: Dict[str, Any]):
|
||||
self.workflow = workflow_json
|
||||
# Create lookups
|
||||
self._nodes = self.workflow
|
||||
if "nodes" in self.workflow and isinstance(self.workflow["nodes"], list):
|
||||
# Handle "graph" format vs "api" format if needed.
|
||||
# But usually for API we stick to the {node_id: node_data} format.
|
||||
# If input is graph format, it might need conversion or distinct handling.
|
||||
# Assuming API format for now as that's what's sent to /prompt.
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def from_json_string(cls, json_str: str) -> 'WorkflowBuilder':
|
||||
"""Load from JSON int."""
|
||||
return cls(json.loads(json_str))
|
||||
|
||||
def get_workflow(self) -> Dict[str, Any]:
|
||||
"""Get the current workflow dict."""
|
||||
return self.workflow
|
||||
|
||||
def set_prompt(self, positive: str, negative: Optional[str] = None) -> None:
|
||||
"""
|
||||
Attempt to set positive and negative prompts.
|
||||
Heuristics:
|
||||
- Look for CLIPTextEncode nodes.
|
||||
- Often one is connected to KSampler 'positive' and one to 'negative'.
|
||||
- Or look for custom titles like 'Positive Prompt', 'Negative Prompt'.
|
||||
"""
|
||||
# Simple heuristic: Find CLIPTextEncode nodes
|
||||
# If we have title/coloring, we can use that.
|
||||
# Otherwise, we might need graph traversal to see what connects to KSampler.
|
||||
|
||||
# For MVP, let's assume standard ComfyUI structure or look for specific titles first
|
||||
|
||||
positive_node_id = self._find_node_by_title("Positive Prompt")
|
||||
negative_node_id = self._find_node_by_title("Negative Prompt")
|
||||
|
||||
# Fallback: Find KSampler and trace back
|
||||
if not positive_node_id or not negative_node_id:
|
||||
ksampler_id, ksampler = self._find_node_by_class("KSampler")
|
||||
if ksampler:
|
||||
# KSampler inputs: model, positive, negative, latent_image
|
||||
if not positive_node_id:
|
||||
positive_node_id = self._trace_input(ksampler, "positive")
|
||||
if not negative_node_id:
|
||||
negative_node_id = self._trace_input(ksampler, "negative")
|
||||
|
||||
if positive_node_id:
|
||||
self._update_node_input(positive_node_id, "text", positive)
|
||||
else:
|
||||
logger.warning("Could not identify Positive Prompt node.")
|
||||
|
||||
if negative and negative_node_id:
|
||||
self._update_node_input(negative_node_id, "text", negative)
|
||||
elif negative:
|
||||
logger.warning("Could not identify Negative Prompt node.")
|
||||
|
||||
def set_seed(self, seed: int) -> int:
|
||||
"""Set seed on KSampler nodes or Seed nodes."""
|
||||
# Find KSampler or anything with a 'seed' widget
|
||||
updated = False
|
||||
for node_id, node in self.workflow.items():
|
||||
if "inputs" in node:
|
||||
if "seed" in node["inputs"]:
|
||||
# Ensure it's an int widget, not a link
|
||||
if isinstance(node["inputs"]["seed"], (int, float)) or (isinstance(node["inputs"]["seed"], str) and node["inputs"]["seed"].isdigit()):
|
||||
node["inputs"]["seed"] = seed
|
||||
updated = True
|
||||
if "noise_seed" in node["inputs"]:
|
||||
# Some nodes call it noise_seed
|
||||
if isinstance(node["inputs"]["noise_seed"], (int, float)):
|
||||
node["inputs"]["noise_seed"] = seed
|
||||
updated = True
|
||||
|
||||
if not updated:
|
||||
logger.warning("Could not find any seed inputs to update.")
|
||||
return seed
|
||||
|
||||
def set_image_dimensions(self, width: int, height: int) -> None:
|
||||
"""Set width and height on EmptyLatentImage nodes."""
|
||||
node_id, _ = self._find_node_by_class("EmptyLatentImage")
|
||||
if node_id:
|
||||
self._update_node_input(node_id, "width", width)
|
||||
self._update_node_input(node_id, "height", height)
|
||||
|
||||
def set_steps(self, steps: int) -> None:
|
||||
"""Set steps on KSampler."""
|
||||
ksampler_ids = self._find_nodes_by_class("KSampler")
|
||||
for nid in ksampler_ids:
|
||||
self._update_node_input(nid, "steps", steps)
|
||||
|
||||
def set_cfg(self, cfg: float) -> None:
|
||||
"""Set CFG scale on KSampler."""
|
||||
ksampler_ids = self._find_nodes_by_class("KSampler")
|
||||
for nid in ksampler_ids:
|
||||
self._update_node_input(nid, "cfg", cfg)
|
||||
|
||||
def _find_node_by_title(self, title: str) -> Optional[str]:
|
||||
"""Find node by its custom title (`_meta.title`)."""
|
||||
for node_id, node in self.workflow.items():
|
||||
if "_meta" in node and node["_meta"].get("title") == title:
|
||||
return node_id
|
||||
return None
|
||||
|
||||
def _find_node_by_class(self, class_type: str) -> Tuple[Optional[str], Optional[Dict]]:
|
||||
"""Find first node of a specific class type."""
|
||||
for node_id, node in self.workflow.items():
|
||||
if node.get("class_type") == class_type:
|
||||
return node_id, node
|
||||
return None, None
|
||||
|
||||
def _find_nodes_by_class(self, class_type: str) -> List[str]:
|
||||
"""Find all nodes of a specific class type."""
|
||||
ids = []
|
||||
for node_id, node in self.workflow.items():
|
||||
if node.get("class_type") == class_type:
|
||||
ids.append(node_id)
|
||||
return ids
|
||||
|
||||
def _trace_input(self, node: Dict, input_name: str) -> Optional[str]:
|
||||
"""
|
||||
Trace back an input link to find the source node.
|
||||
Input format in API JSON: "input_name": ["source_node_id", slot_index]
|
||||
"""
|
||||
if "inputs" not in node or input_name not in node["inputs"]:
|
||||
return None
|
||||
|
||||
link = node["inputs"][input_name]
|
||||
# Link structure: [node_id, slot_idx]
|
||||
if isinstance(link, list) and len(link) == 2:
|
||||
return str(link[0])
|
||||
return None
|
||||
|
||||
def _update_node_input(self, node_id: str, input_name: str, value: Any) -> None:
|
||||
if node_id in self.workflow and "inputs" in self.workflow[node_id]:
|
||||
self.workflow[node_id]["inputs"][input_name] = value
|
||||
@@ -0,0 +1,48 @@
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Add project root to python path
|
||||
project_root = Path(__file__).resolve().parent.parent.parent
|
||||
sys.path.append(str(project_root))
|
||||
|
||||
print(f"Testing imports from {project_root}")
|
||||
|
||||
try:
|
||||
import bot.config
|
||||
print("✅ bot.config imported")
|
||||
|
||||
import bot.database.models
|
||||
print("✅ bot.database.models imported")
|
||||
|
||||
import bot.database.repository
|
||||
print("✅ bot.database.repository imported")
|
||||
|
||||
import bot.comfyui.client
|
||||
print("✅ bot.comfyui.client imported")
|
||||
|
||||
import bot.services.job_manager
|
||||
print("✅ bot.services.job_manager imported")
|
||||
|
||||
import bot.cogs.generate
|
||||
print("✅ bot.cogs.generate imported")
|
||||
|
||||
import bot.services.permissions
|
||||
print("✅ bot.services.permissions imported")
|
||||
|
||||
import bot.cogs.queue
|
||||
print("✅ bot.cogs.queue imported")
|
||||
|
||||
import bot.cogs.admin
|
||||
print("✅ bot.cogs.admin imported")
|
||||
|
||||
import bot.bot
|
||||
print("✅ bot.bot imported")
|
||||
|
||||
print("All modules imported successfully.")
|
||||
except ImportError as e:
|
||||
print(f"❌ ImportError: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
sys.exit(1)
|
||||
@@ -0,0 +1,52 @@
|
||||
"""
|
||||
Pytest configuration file for tests.
|
||||
|
||||
This module handles test isolation by ensuring that real packages are imported
|
||||
before any mocking occurs, and by providing cleanup fixtures.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add project root to path
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# Store references to real modules before any mocking
|
||||
# This ensures tests that need real numpy/PIL can use them
|
||||
_real_numpy = None
|
||||
_real_PIL = None
|
||||
|
||||
def pytest_configure(config):
|
||||
"""Called after command line options have been parsed and all plugins loaded."""
|
||||
global _real_numpy, _real_PIL
|
||||
|
||||
# Import real modules and store references
|
||||
try:
|
||||
import numpy
|
||||
_real_numpy = numpy
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import PIL
|
||||
import PIL.Image
|
||||
import PIL.PngImagePlugin
|
||||
_real_PIL = PIL
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def get_real_numpy():
|
||||
"""Get the real numpy module, not a mock."""
|
||||
if _real_numpy is None:
|
||||
import numpy
|
||||
return numpy
|
||||
return _real_numpy
|
||||
|
||||
|
||||
def get_real_PIL():
|
||||
"""Get the real PIL module, not a mock."""
|
||||
if _real_PIL is None:
|
||||
import PIL
|
||||
return PIL
|
||||
return _real_PIL
|
||||
@@ -0,0 +1,94 @@
|
||||
import unittest
|
||||
import sys
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.discord import send_to_discord_with_retry
|
||||
|
||||
class TestDiscordAPI(unittest.TestCase):
|
||||
"""Tests for the Discord API utility with mocked network responses."""
|
||||
|
||||
@patch('requests.post')
|
||||
@patch('time.sleep', return_value=None) # Don't actually wait during tests
|
||||
def test_retry_on_rate_limit(self, mock_sleep, mock_post):
|
||||
"""Should retry when receiving a 429 Rate Limit error."""
|
||||
# Setup mock responses: first is 429, second is 200
|
||||
mock_429 = MagicMock()
|
||||
mock_429.status_code = 429
|
||||
mock_429.text = "Too many requests"
|
||||
mock_429.json.return_value = {"retry_after": 0.1}
|
||||
|
||||
mock_200 = MagicMock()
|
||||
mock_200.status_code = 200
|
||||
|
||||
mock_post.side_effect = [mock_429, mock_200]
|
||||
|
||||
response = send_to_discord_with_retry(
|
||||
"https://discord.com/api/webhooks/123/abc",
|
||||
data={"content": "test message"},
|
||||
max_retries=2
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
self.assertEqual(mock_post.call_count, 2)
|
||||
mock_sleep.assert_called()
|
||||
|
||||
@patch('requests.post')
|
||||
def test_no_retry_on_400(self, mock_post):
|
||||
"""Should not retry and return response for 400 Bad Request."""
|
||||
mock_400 = MagicMock()
|
||||
mock_400.status_code = 400
|
||||
mock_400.text = "Bad Request"
|
||||
mock_post.return_value = mock_400
|
||||
|
||||
response = send_to_discord_with_retry(
|
||||
"https://discord.com/api/webhooks/123/abc",
|
||||
data={"content": "test message"}
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 400)
|
||||
self.assertEqual(mock_post.call_count, 1)
|
||||
|
||||
@patch('requests.post')
|
||||
@patch('time.sleep', return_value=None)
|
||||
def test_exhaust_retries(self, mock_sleep, mock_post):
|
||||
"""Should return the last failed response after exhaust all retries."""
|
||||
mock_500 = MagicMock()
|
||||
mock_500.status_code = 500
|
||||
mock_post.return_value = mock_500
|
||||
|
||||
# Should return the last 500 response
|
||||
response = send_to_discord_with_retry(
|
||||
"https://discord.com/api/webhooks/131072/abc", # Valid ID
|
||||
data={"content": "test message"},
|
||||
max_retries=3
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 500)
|
||||
self.assertEqual(mock_post.call_count, 3)
|
||||
self.assertEqual(mock_sleep.call_count, 2) # Sleeps between attempts (1st to 2nd, 2nd to 3rd)
|
||||
|
||||
@patch('requests.post')
|
||||
@patch('time.sleep', return_value=None)
|
||||
def test_raise_exception_on_network_failure(self, mock_sleep, mock_post):
|
||||
"""Should raise RequestException if all retries hit network errors."""
|
||||
import requests
|
||||
mock_post.side_effect = requests.exceptions.ConnectionError("Failed")
|
||||
|
||||
with self.assertRaises(requests.exceptions.ConnectionError):
|
||||
send_to_discord_with_retry(
|
||||
"https://discord.com/api/webhooks/123/abc",
|
||||
data={"content": "test"},
|
||||
max_retries=2
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Tests for shared/filename_utils.py"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.filename_utils import build_filename_with_metadata, get_timestamp_string
|
||||
|
||||
|
||||
class TestBuildFilenameWithMetadata(unittest.TestCase):
|
||||
"""Test build_filename_with_metadata function."""
|
||||
|
||||
def test_prefix_only(self):
|
||||
"""Test with just a prefix, no metadata."""
|
||||
result, info = build_filename_with_metadata("image")
|
||||
self.assertEqual(result, "image")
|
||||
self.assertEqual(info, {})
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_with_date(self, mock_time):
|
||||
"""Test adding date to filename."""
|
||||
mock_time.strftime.return_value = "2026-01-20"
|
||||
result, info = build_filename_with_metadata("image", add_date=True)
|
||||
self.assertEqual(result, "image_2026-01-20")
|
||||
self.assertEqual(info["date"], "2026-01-20")
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_with_time(self, mock_time):
|
||||
"""Test adding time to filename."""
|
||||
mock_time.strftime.return_value = "14-30-00"
|
||||
result, info = build_filename_with_metadata("image", add_time=True)
|
||||
self.assertEqual(result, "image_14-30-00")
|
||||
self.assertEqual(info["time"], "14-30-00")
|
||||
|
||||
def test_with_dimensions(self):
|
||||
"""Test adding dimensions to filename."""
|
||||
result, info = build_filename_with_metadata(
|
||||
"image", add_dimensions=True, width=1920, height=1080
|
||||
)
|
||||
self.assertEqual(result, "image_1920x1080")
|
||||
self.assertEqual(info["dimensions"], "1920x1080")
|
||||
|
||||
def test_dimensions_without_values(self):
|
||||
"""Test that dimensions are not added without width/height."""
|
||||
result, info = build_filename_with_metadata("image", add_dimensions=True)
|
||||
self.assertEqual(result, "image")
|
||||
self.assertNotIn("dimensions", info)
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_all_metadata(self, mock_time):
|
||||
"""Test with all metadata options."""
|
||||
mock_time.strftime.side_effect = ["2026-01-20", "14-30-00"]
|
||||
result, info = build_filename_with_metadata(
|
||||
"output",
|
||||
add_date=True,
|
||||
add_time=True,
|
||||
add_dimensions=True,
|
||||
width=512,
|
||||
height=768,
|
||||
)
|
||||
self.assertEqual(result, "output_2026-01-20_14-30-00_512x768")
|
||||
self.assertEqual(info["date"], "2026-01-20")
|
||||
self.assertEqual(info["time"], "14-30-00")
|
||||
self.assertEqual(info["dimensions"], "512x768")
|
||||
|
||||
def test_with_existing_info_dict(self):
|
||||
"""Test that existing info_dict is updated, not replaced."""
|
||||
existing_info = {"existing_key": "existing_value"}
|
||||
result, info = build_filename_with_metadata(
|
||||
"image", add_dimensions=True, width=100, height=100, info_dict=existing_info
|
||||
)
|
||||
self.assertEqual(info["existing_key"], "existing_value")
|
||||
self.assertEqual(info["dimensions"], "100x100")
|
||||
self.assertIs(info, existing_info) # Same dict object
|
||||
|
||||
|
||||
class TestGetTimestampString(unittest.TestCase):
|
||||
"""Test get_timestamp_string function."""
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_date_only(self, mock_time):
|
||||
"""Test timestamp with date only."""
|
||||
mock_time.strftime.return_value = "2026-01-20"
|
||||
result = get_timestamp_string(include_date=True, include_time=False)
|
||||
self.assertEqual(result, "2026-01-20")
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_time_only(self, mock_time):
|
||||
"""Test timestamp with time only."""
|
||||
mock_time.strftime.return_value = "14-30-00"
|
||||
result = get_timestamp_string(include_date=False, include_time=True)
|
||||
self.assertEqual(result, "14-30-00")
|
||||
|
||||
@patch("shared.filename_utils.time")
|
||||
def test_both(self, mock_time):
|
||||
"""Test timestamp with both date and time."""
|
||||
mock_time.strftime.side_effect = ["2026-01-20", "14-30-00"]
|
||||
result = get_timestamp_string(include_date=True, include_time=True)
|
||||
self.assertEqual(result, "2026-01-20_14-30-00")
|
||||
|
||||
def test_neither(self):
|
||||
"""Test timestamp with neither date nor time."""
|
||||
result = get_timestamp_string(include_date=False, include_time=False)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,216 @@
|
||||
"""Tests for shared/media/format_utils.py"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
import tempfile
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.media.format_utils import (
|
||||
parse_format_string,
|
||||
normalize_video_extension,
|
||||
get_mime_type,
|
||||
validate_video_for_discord,
|
||||
is_animated_format,
|
||||
supports_alpha,
|
||||
)
|
||||
|
||||
|
||||
class TestParseFormatString(unittest.TestCase):
|
||||
"""Test parse_format_string function."""
|
||||
|
||||
def test_video_h264_mp4(self):
|
||||
"""Test parsing video/h264-mp4 format."""
|
||||
fmt_type, fmt_ext = parse_format_string("video/h264-mp4")
|
||||
self.assertEqual(fmt_type, "video")
|
||||
self.assertEqual(fmt_ext, "h264-mp4")
|
||||
|
||||
def test_image_gif(self):
|
||||
"""Test parsing image/gif format."""
|
||||
fmt_type, fmt_ext = parse_format_string("image/gif")
|
||||
self.assertEqual(fmt_type, "image")
|
||||
self.assertEqual(fmt_ext, "gif")
|
||||
|
||||
def test_simple_format(self):
|
||||
"""Test parsing simple format string without slash."""
|
||||
fmt_type, fmt_ext = parse_format_string("mp4")
|
||||
self.assertEqual(fmt_type, "video")
|
||||
self.assertEqual(fmt_ext, "mp4")
|
||||
|
||||
|
||||
class TestNormalizeVideoExtension(unittest.TestCase):
|
||||
"""Test normalize_video_extension function."""
|
||||
|
||||
def test_h264_mp4(self):
|
||||
"""Test normalizing h264-mp4 to mp4."""
|
||||
self.assertEqual(normalize_video_extension("video/h264-mp4"), "mp4")
|
||||
|
||||
def test_h265_mp4(self):
|
||||
"""Test normalizing h265-mp4 to mp4."""
|
||||
self.assertEqual(normalize_video_extension("video/h265-mp4"), "mp4")
|
||||
|
||||
def test_vp9_webm(self):
|
||||
"""Test normalizing vp9-webm to webm."""
|
||||
self.assertEqual(normalize_video_extension("video/vp9-webm"), "webm")
|
||||
|
||||
def test_prores(self):
|
||||
"""Test normalizing prores to mov."""
|
||||
self.assertEqual(normalize_video_extension("video/prores"), "mov")
|
||||
|
||||
def test_gif_passthrough(self):
|
||||
"""Test gif format passes through unchanged."""
|
||||
self.assertEqual(normalize_video_extension("image/gif"), "gif")
|
||||
|
||||
def test_unknown_passthrough(self):
|
||||
"""Test unknown format passes through unchanged."""
|
||||
self.assertEqual(normalize_video_extension("video/custom"), "custom")
|
||||
|
||||
|
||||
class TestGetMimeType(unittest.TestCase):
|
||||
"""Test get_mime_type function."""
|
||||
|
||||
def test_mp4(self):
|
||||
"""Test MIME type for mp4."""
|
||||
self.assertEqual(get_mime_type("mp4"), "video/mp4")
|
||||
|
||||
def test_webm(self):
|
||||
"""Test MIME type for webm."""
|
||||
self.assertEqual(get_mime_type("webm"), "video/webm")
|
||||
|
||||
def test_gif(self):
|
||||
"""Test MIME type for gif."""
|
||||
self.assertEqual(get_mime_type("gif"), "image/gif")
|
||||
|
||||
def test_mov(self):
|
||||
"""Test MIME type for mov."""
|
||||
self.assertEqual(get_mime_type("mov"), "video/quicktime")
|
||||
|
||||
def test_case_insensitive(self):
|
||||
"""Test MIME type lookup is case insensitive."""
|
||||
self.assertEqual(get_mime_type("MP4"), "video/mp4")
|
||||
|
||||
def test_unknown_format(self):
|
||||
"""Test unknown format returns octet-stream."""
|
||||
self.assertEqual(get_mime_type("xyz"), "application/octet-stream")
|
||||
|
||||
|
||||
class TestValidateVideoForDiscord(unittest.TestCase):
|
||||
"""Test validate_video_for_discord function."""
|
||||
|
||||
def test_nonexistent_file(self):
|
||||
"""Test validation of nonexistent file."""
|
||||
is_valid, msg = validate_video_for_discord("/nonexistent/file.mp4")
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("does not exist", msg)
|
||||
|
||||
def test_empty_file(self):
|
||||
"""Test validation of empty file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
|
||||
temp_path = f.name
|
||||
try:
|
||||
is_valid, msg = validate_video_for_discord(temp_path)
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("empty", msg)
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_small_file(self):
|
||||
"""Test validation of suspiciously small file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
|
||||
f.write(b"x" * 100) # 100 bytes
|
||||
temp_path = f.name
|
||||
try:
|
||||
is_valid, msg = validate_video_for_discord(temp_path)
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("small", msg)
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_valid_mp4(self):
|
||||
"""Test validation of valid mp4 file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
|
||||
f.write(b"x" * 10000) # 10KB
|
||||
temp_path = f.name
|
||||
try:
|
||||
is_valid, msg = validate_video_for_discord(temp_path)
|
||||
self.assertTrue(is_valid)
|
||||
self.assertEqual(msg, "Valid")
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_valid_webm(self):
|
||||
"""Test validation of valid webm file."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".webm", delete=False) as f:
|
||||
f.write(b"x" * 10000)
|
||||
temp_path = f.name
|
||||
try:
|
||||
is_valid, msg = validate_video_for_discord(temp_path)
|
||||
self.assertTrue(is_valid)
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
def test_mov_needs_conversion(self):
|
||||
"""Test that MOV files are flagged for conversion."""
|
||||
with tempfile.NamedTemporaryFile(suffix=".mov", delete=False) as f:
|
||||
f.write(b"x" * 10000)
|
||||
temp_path = f.name
|
||||
try:
|
||||
is_valid, msg = validate_video_for_discord(temp_path)
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("conversion", msg)
|
||||
finally:
|
||||
os.unlink(temp_path)
|
||||
|
||||
|
||||
class TestIsAnimatedFormat(unittest.TestCase):
|
||||
"""Test is_animated_format function."""
|
||||
|
||||
def test_animated_formats(self):
|
||||
"""Test formats that support animation."""
|
||||
animated = ["gif", "webp", "mp4", "webm", "mov", "avi", "mkv", "apng"]
|
||||
for fmt in animated:
|
||||
self.assertTrue(is_animated_format(fmt), f"{fmt} should be animated")
|
||||
|
||||
def test_static_formats(self):
|
||||
"""Test formats that don't support animation."""
|
||||
static = ["png", "jpg", "jpeg", "bmp"]
|
||||
for fmt in static:
|
||||
self.assertFalse(is_animated_format(fmt), f"{fmt} should not be animated")
|
||||
|
||||
def test_case_insensitive(self):
|
||||
"""Test case insensitivity."""
|
||||
self.assertTrue(is_animated_format("GIF"))
|
||||
self.assertTrue(is_animated_format("Mp4"))
|
||||
|
||||
|
||||
class TestSupportsAlpha(unittest.TestCase):
|
||||
"""Test supports_alpha function."""
|
||||
|
||||
def test_alpha_formats(self):
|
||||
"""Test formats that support alpha channel."""
|
||||
alpha = ["webm", "gif", "webp", "png", "apng", "mov"]
|
||||
for fmt in alpha:
|
||||
self.assertTrue(supports_alpha(fmt), f"{fmt} should support alpha")
|
||||
|
||||
def test_no_alpha_formats(self):
|
||||
"""Test formats that don't support alpha."""
|
||||
no_alpha = ["mp4", "jpg", "jpeg", "avi"]
|
||||
for fmt in no_alpha:
|
||||
self.assertFalse(supports_alpha(fmt), f"{fmt} should not support alpha")
|
||||
|
||||
def test_case_insensitive(self):
|
||||
"""Test case insensitivity."""
|
||||
self.assertTrue(supports_alpha("PNG"))
|
||||
self.assertTrue(supports_alpha("WebM"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,99 @@
|
||||
|
||||
import unittest
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock torch and other heavy dependencies
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["PIL"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
sys.modules["folder_paths"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["comfy.cli_args"] = MagicMock()
|
||||
|
||||
# Now we can safely import
|
||||
from shared.github_integration import validate_github_repo, validate_file_path, update_github_cdn_urls
|
||||
|
||||
class TestGitHubValidation(unittest.TestCase):
|
||||
|
||||
def test_validate_github_repo_valid(self):
|
||||
"""Test valid GitHub repository formats."""
|
||||
valid_repos = [
|
||||
"username/repo",
|
||||
"user-name/repo-name",
|
||||
"user-name/repo.name", # Dot in repo is valid
|
||||
"user-name/repo_name", # Underscore in repo is valid
|
||||
"0123/4567"
|
||||
]
|
||||
for repo in valid_repos:
|
||||
with self.subTest(repo=repo):
|
||||
self.assertTrue(validate_github_repo(repo), f"Failed for {repo}")
|
||||
|
||||
def test_validate_github_repo_invalid(self):
|
||||
"""Test invalid GitHub repository formats (traversal, injection, invalid chars)."""
|
||||
invalid_repos = [
|
||||
"username/repo/../other", # Traversal
|
||||
"username/repo?query=1", # Query injection
|
||||
"username", # Missing slash
|
||||
"/repo", # Missing username
|
||||
"user/", # Missing repo
|
||||
"user/repo/", # Trailing slash (strict check)
|
||||
"../../user/repo", # Traversal at start
|
||||
"user/repo#fragment", # Fragment
|
||||
"user/repo;rm -rf", # Command injection style
|
||||
"user/repo\nnewline", # Newline
|
||||
"user.name/repo", # Dot in username (invalid)
|
||||
"user_name/repo", # Underscore in username (invalid)
|
||||
]
|
||||
for repo in invalid_repos:
|
||||
with self.subTest(repo=repo):
|
||||
self.assertFalse(validate_github_repo(repo), f"Should have failed for {repo}")
|
||||
|
||||
def test_validate_file_path_valid(self):
|
||||
"""Test valid file paths."""
|
||||
valid_paths = [
|
||||
"file.txt",
|
||||
"path/to/file.txt",
|
||||
"folder/subfolder/file.md",
|
||||
"README.md",
|
||||
"docs/image.png"
|
||||
]
|
||||
for path in valid_paths:
|
||||
with self.subTest(path=path):
|
||||
self.assertTrue(validate_file_path(path), f"Failed for {path}")
|
||||
|
||||
def test_validate_file_path_invalid(self):
|
||||
"""Test invalid file paths (traversal, absolute)."""
|
||||
invalid_paths = [
|
||||
"../file.txt", # Traversal
|
||||
"path/../file.txt", # Traversal inside
|
||||
"/etc/passwd", # Absolute path
|
||||
"/file.txt", # Absolute path
|
||||
"../../secret", # Deep traversal
|
||||
"", # Empty
|
||||
None # None
|
||||
]
|
||||
for path in invalid_paths:
|
||||
with self.subTest(path=path):
|
||||
self.assertFalse(validate_file_path(path), f"Should have failed for {path}")
|
||||
|
||||
def test_update_github_cdn_urls_rejects_invalid_repo(self):
|
||||
"""Test that update_github_cdn_urls rejects invalid repo before making requests."""
|
||||
repo = "user/repo/../malicious"
|
||||
success, message = update_github_cdn_urls(repo, "token", "file.md", [("f", "u")])
|
||||
|
||||
self.assertFalse(success)
|
||||
self.assertIn("Invalid GitHub repository format", message)
|
||||
|
||||
def test_update_github_cdn_urls_rejects_invalid_path(self):
|
||||
"""Test that update_github_cdn_urls rejects invalid path before making requests."""
|
||||
path = "../../../secret.txt"
|
||||
success, message = update_github_cdn_urls("user/repo", "token", path, [("f", "u")])
|
||||
|
||||
self.assertFalse(success)
|
||||
self.assertIn("Invalid file path", message)
|
||||
self.assertIn("Path traversal", message)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,173 @@
|
||||
import unittest
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
from unittest.mock import MagicMock, patch
|
||||
import importlib
|
||||
|
||||
# Add project root to sys.path
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
# Clean up any potential pollution from other tests before we start
|
||||
if 'PIL' in sys.modules:
|
||||
# Check if it's a mock
|
||||
if isinstance(sys.modules['PIL'], MagicMock):
|
||||
del sys.modules['PIL']
|
||||
if 'PIL.PngImagePlugin' in sys.modules:
|
||||
del sys.modules['PIL.PngImagePlugin']
|
||||
if 'PIL.Image' in sys.modules:
|
||||
del sys.modules['PIL.Image']
|
||||
|
||||
# Now we can import real modules or mock them as we see fit LOCALLY
|
||||
# But wait, discord_image_node imports them at module level.
|
||||
# So we need to ensure environment is set up before importing it.
|
||||
|
||||
# Mock comfy modules
|
||||
sys.modules['comfy'] = MagicMock()
|
||||
sys.modules['comfy.cli_args'] = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args.disable_metadata = False
|
||||
sys.modules['comfy.utils'] = MagicMock()
|
||||
sys.modules['folder_paths'] = MagicMock()
|
||||
sys.modules['folder_paths'].get_output_directory = MagicMock(return_value="/tmp")
|
||||
sys.modules['folder_paths'].get_temp_directory = MagicMock(return_value="/tmp")
|
||||
sys.modules['folder_paths'].get_save_image_path = MagicMock(return_value=("/tmp", "test", 0, "", "test"))
|
||||
sys.modules['server'] = MagicMock()
|
||||
|
||||
# We need real PIL for this test to verify PngInfo
|
||||
try:
|
||||
import PIL.PngImagePlugin
|
||||
except ImportError:
|
||||
# If it failed because it was mocked out and we deleted it, reload
|
||||
pass
|
||||
|
||||
from nodes.image_node import DiscordSendSaveImage
|
||||
|
||||
# Check if torch is real or mocked
|
||||
try:
|
||||
import torch
|
||||
_torch_available = hasattr(torch, 'zeros') and callable(torch.zeros) and not isinstance(torch.zeros, MagicMock)
|
||||
except ImportError:
|
||||
_torch_available = False
|
||||
|
||||
class TestDiscordImageNodeOptimization(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.node = DiscordSendSaveImage()
|
||||
self.webhook_url = "https://discord.com/api/webhooks/12345/abcdef"
|
||||
self.github_token = "ghp_sensitive12345"
|
||||
|
||||
@unittest.skipUnless(_torch_available, "Test requires real torch for tensor iteration")
|
||||
def test_save_images_sanitization(self):
|
||||
# Create a mock image tensor using numpy (torch not available in CI)
|
||||
# The image_node iterates over images and accesses shape, so we need
|
||||
# an object that supports iteration and has proper shape
|
||||
import numpy as np
|
||||
|
||||
# Create a simple class that mimics torch.Tensor behavior for the node
|
||||
class MockTensor:
|
||||
def __init__(self, data):
|
||||
self._data = data
|
||||
self.shape = data.shape
|
||||
|
||||
def __len__(self):
|
||||
return len(self._data)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self._data[idx]
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._data)
|
||||
|
||||
# Create a 1x64x64x3 "image batch" using numpy
|
||||
image_data = np.zeros((1, 64, 64, 3), dtype=np.float32)
|
||||
mock_image = MockTensor(image_data)
|
||||
|
||||
# Create prompt and extra_pnginfo with sensitive data
|
||||
prompt = {
|
||||
"3": {
|
||||
"inputs": {
|
||||
"webhook_url": self.webhook_url,
|
||||
"github_token": self.github_token,
|
||||
"seed": 123
|
||||
},
|
||||
"class_type": "DiscordSendSaveImage"
|
||||
}
|
||||
}
|
||||
|
||||
extra_pnginfo = {
|
||||
"workflow": {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 3,
|
||||
"type": "DiscordSendSaveImage",
|
||||
"widgets_values": [self.webhook_url, "message", self.github_token]
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
# Mock Image.save to check metadata
|
||||
# Also mock tensor_to_numpy_uint8 to bypass torch tensor conversion (torch is mocked)
|
||||
# and Image.fromarray to return a mock PIL Image with proper size attribute
|
||||
mock_pil_image = MagicMock()
|
||||
mock_pil_image.size = (64, 64)
|
||||
mock_pil_image.mode = 'RGB'
|
||||
|
||||
def mock_tensor_to_numpy(tensor):
|
||||
# Return a simple numpy-like array (64x64x3 zeros as uint8)
|
||||
import numpy as np
|
||||
return np.zeros((64, 64, 3), dtype=np.uint8)
|
||||
|
||||
with patch('PIL.Image.Image.save') as mock_save, \
|
||||
patch('nodes.image_node.tensor_to_numpy_uint8', side_effect=mock_tensor_to_numpy), \
|
||||
patch('PIL.Image.fromarray', return_value=mock_pil_image):
|
||||
self.node.save_images(
|
||||
images=mock_image,
|
||||
prompt=prompt,
|
||||
extra_pnginfo=extra_pnginfo,
|
||||
save_output=True,
|
||||
send_to_discord=False # Disable discord sending to focus on save/metadata
|
||||
)
|
||||
|
||||
# Check if save was called
|
||||
self.assertTrue(mock_save.called)
|
||||
|
||||
# Get the pnginfo passed to save
|
||||
args, kwargs = mock_save.call_args
|
||||
pnginfo = kwargs.get('pnginfo')
|
||||
self.assertIsNotNone(pnginfo)
|
||||
|
||||
found_prompt = False
|
||||
found_workflow = False
|
||||
|
||||
# Check chunks - PIL PngInfo internal structure
|
||||
for tag_type, data, after_idat in pnginfo.chunks:
|
||||
try:
|
||||
# decode data
|
||||
decoded = data.decode('latin-1')
|
||||
except Exception:
|
||||
# Skip chunks that can't be decoded
|
||||
continue
|
||||
|
||||
if '\0' in decoded:
|
||||
try:
|
||||
k, v = decoded.split('\0', 1)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if k == "prompt":
|
||||
found_prompt = True
|
||||
# Verify sensitive data is gone
|
||||
self.assertNotIn("discord.com/api/webhooks", v)
|
||||
self.assertNotIn("ghp_", v)
|
||||
|
||||
if k == "workflow":
|
||||
found_workflow = True
|
||||
self.assertNotIn("discord.com/api/webhooks", v)
|
||||
self.assertNotIn("ghp_", v)
|
||||
|
||||
self.assertTrue(found_prompt, "Prompt metadata not found")
|
||||
self.assertTrue(found_workflow, "Workflow metadata not found")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,85 @@
|
||||
# Mock dependencies before imports
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# IMPORTANT: Import real numpy FIRST before any mocking
|
||||
# This ensures test_power_of_two_math uses real numpy
|
||||
import numpy as np
|
||||
|
||||
# Create a dummy torch module (torch is not installed in CI)
|
||||
mock_torch = MagicMock()
|
||||
sys.modules["torch"] = mock_torch
|
||||
sys.modules["folder_paths"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["comfy.cli_args"] = MagicMock()
|
||||
sys.modules["comfy.utils"] = MagicMock()
|
||||
sys.modules["PIL"] = MagicMock()
|
||||
sys.modules["PIL.Image"] = MagicMock()
|
||||
sys.modules["PIL.PngImagePlugin"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
sys.modules["server"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
import os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from nodes.video_node import validate_video_for_discord
|
||||
|
||||
class TestPathLogic(unittest.TestCase):
|
||||
"""Tests for path and file validation logic."""
|
||||
|
||||
@patch('os.path.exists')
|
||||
@patch('os.path.getsize')
|
||||
def test_validate_video_discord_limits(self, mock_getsize, mock_exists):
|
||||
"""Should correctly identify if a video is within Discord limits."""
|
||||
mock_exists.return_value = True
|
||||
|
||||
# Test 10MB file (OK)
|
||||
mock_getsize.return_value = 10 * 1024 * 1024
|
||||
is_valid, msg = validate_video_for_discord("test.mp4")
|
||||
self.assertTrue(is_valid)
|
||||
|
||||
# Test 100MB file (Too large)
|
||||
mock_getsize.return_value = 100 * 1024 * 1024
|
||||
is_valid, msg = validate_video_for_discord("test.mp4")
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("exceeds Discord's size limit", msg)
|
||||
|
||||
def test_validate_video_formats(self):
|
||||
"""Should validate supported video formats."""
|
||||
with patch('os.path.exists', return_value=True), \
|
||||
patch('os.path.getsize', return_value=1024 * 1024):
|
||||
|
||||
self.assertTrue(validate_video_for_discord("test.mp4")[0])
|
||||
self.assertTrue(validate_video_for_discord("test.webm")[0])
|
||||
self.assertTrue(validate_video_for_discord("test.gif")[0])
|
||||
|
||||
# MOV usually marked as needing conversion
|
||||
is_valid, msg = validate_video_for_discord("test.mov")
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("need conversion", msg.lower())
|
||||
|
||||
class TestImageResizing(unittest.TestCase):
|
||||
"""
|
||||
Tests for image resizing logic.
|
||||
Note: The actual resizing logic is inside the Node class.
|
||||
We are testing the mathematical logic that would be used.
|
||||
"""
|
||||
|
||||
def test_power_of_two_math(self):
|
||||
"""Verify the power-of-two calculation logic used in the node."""
|
||||
# Use Python's built-in math module instead of numpy
|
||||
# to avoid test collection order issues with mocked modules
|
||||
import math
|
||||
|
||||
def calculate_nearest_pow2(dim):
|
||||
return 2 ** int(math.log2(dim) + 0.5)
|
||||
|
||||
self.assertEqual(calculate_nearest_pow2(500), 512)
|
||||
self.assertEqual(calculate_nearest_pow2(700), 512) # log2(700) = 9.45, +0.5 = 9.95, int=9, 2^9=512
|
||||
self.assertEqual(calculate_nearest_pow2(800), 1024) # log2(800) = 9.64, +0.5 = 10.14, int=10, 2^10=1024
|
||||
self.assertEqual(calculate_nearest_pow2(256), 256)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,261 @@
|
||||
"""Tests for shared/discord/message_builder.py"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.discord.message_builder import (
|
||||
build_metadata_section,
|
||||
build_prompt_section,
|
||||
build_discord_message,
|
||||
validate_message_content,
|
||||
format_file_info,
|
||||
format_file_size,
|
||||
)
|
||||
|
||||
|
||||
class TestBuildMetadataSection(unittest.TestCase):
|
||||
"""Test build_metadata_section function."""
|
||||
|
||||
def test_empty_dict(self):
|
||||
"""Test with empty info dict returns empty string."""
|
||||
result = build_metadata_section({})
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_with_date(self):
|
||||
"""Test metadata with date."""
|
||||
result = build_metadata_section({"date": "2026-01-20"})
|
||||
self.assertIn("**Date:** 2026-01-20", result)
|
||||
self.assertIn("**Information:**", result)
|
||||
|
||||
def test_with_time(self):
|
||||
"""Test metadata with time."""
|
||||
result = build_metadata_section({"time": "14-30-00"})
|
||||
self.assertIn("**Time:** 14-30-00", result)
|
||||
|
||||
def test_with_dimensions(self):
|
||||
"""Test metadata with dimensions."""
|
||||
result = build_metadata_section({"dimensions": "1920x1080"})
|
||||
self.assertIn("**Dimensions:** 1920x1080", result)
|
||||
|
||||
def test_with_format(self):
|
||||
"""Test metadata with file format."""
|
||||
result = build_metadata_section({}, file_format="png")
|
||||
self.assertIn("**Format:** PNG", result)
|
||||
|
||||
def test_with_frame_rate(self):
|
||||
"""Test metadata with frame rate."""
|
||||
result = build_metadata_section({}, frame_rate=30.0)
|
||||
self.assertIn("**Frame Rate:** 30.0 fps", result)
|
||||
|
||||
def test_custom_section_title(self):
|
||||
"""Test custom section title."""
|
||||
result = build_metadata_section({"date": "2026-01-20"}, section_title="Video Info")
|
||||
self.assertIn("**Video Info:**", result)
|
||||
|
||||
def test_exclude_options(self):
|
||||
"""Test excluding certain metadata."""
|
||||
info = {"date": "2026-01-20", "time": "14-30-00", "dimensions": "512x512"}
|
||||
result = build_metadata_section(info, include_date=False, include_time=False)
|
||||
self.assertNotIn("Date", result)
|
||||
self.assertNotIn("Time", result)
|
||||
self.assertIn("Dimensions", result)
|
||||
|
||||
def test_trailing_newline(self):
|
||||
"""Test that section ends with newline."""
|
||||
result = build_metadata_section({"date": "2026-01-20"})
|
||||
self.assertTrue(result.endswith("\n"))
|
||||
|
||||
|
||||
class TestBuildPromptSection(unittest.TestCase):
|
||||
"""Test build_prompt_section function."""
|
||||
|
||||
def test_no_prompts(self):
|
||||
"""Test with no prompts returns empty string."""
|
||||
result = build_prompt_section(None, None)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_empty_prompts(self):
|
||||
"""Test with empty prompts returns empty string."""
|
||||
result = build_prompt_section("", "")
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_whitespace_prompts(self):
|
||||
"""Test with whitespace-only prompts returns empty string."""
|
||||
result = build_prompt_section(" ", " \n ")
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_positive_only(self):
|
||||
"""Test with only positive prompt."""
|
||||
result = build_prompt_section("a beautiful sunset", None)
|
||||
self.assertIn("**Positive:**", result)
|
||||
self.assertIn("a beautiful sunset", result)
|
||||
self.assertNotIn("**Negative:**", result)
|
||||
|
||||
def test_negative_only(self):
|
||||
"""Test with only negative prompt."""
|
||||
result = build_prompt_section(None, "blurry, low quality")
|
||||
self.assertIn("**Negative:**", result)
|
||||
self.assertIn("blurry, low quality", result)
|
||||
self.assertNotIn("**Positive:**", result)
|
||||
|
||||
def test_both_prompts(self):
|
||||
"""Test with both prompts."""
|
||||
result = build_prompt_section("a cat", "dog")
|
||||
self.assertIn("**Positive:**", result)
|
||||
self.assertIn("a cat", result)
|
||||
self.assertIn("**Negative:**", result)
|
||||
self.assertIn("dog", result)
|
||||
|
||||
def test_custom_section_title(self):
|
||||
"""Test custom section title."""
|
||||
result = build_prompt_section("test", None, section_title="Custom Prompts")
|
||||
self.assertIn("**Custom Prompts:**", result)
|
||||
|
||||
def test_code_block_formatting(self):
|
||||
"""Test prompts are wrapped in code blocks."""
|
||||
result = build_prompt_section("test prompt", None)
|
||||
self.assertIn("```\ntest prompt\n```", result)
|
||||
|
||||
def test_non_string_conversion(self):
|
||||
"""Test that non-string prompts are converted."""
|
||||
result = build_prompt_section(12345, None)
|
||||
self.assertIn("12345", result)
|
||||
|
||||
|
||||
class TestBuildDiscordMessage(unittest.TestCase):
|
||||
"""Test build_discord_message function."""
|
||||
|
||||
def test_empty_message(self):
|
||||
"""Test building empty message."""
|
||||
result = build_discord_message()
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_base_message_only(self):
|
||||
"""Test with just base message."""
|
||||
result = build_discord_message(base_message="Hello!")
|
||||
self.assertEqual(result, "Hello!")
|
||||
|
||||
def test_with_metadata(self):
|
||||
"""Test with metadata section."""
|
||||
result = build_discord_message(
|
||||
base_message="Image generated",
|
||||
metadata_section="\n**Info:** test"
|
||||
)
|
||||
self.assertIn("Image generated", result)
|
||||
self.assertIn("**Info:** test", result)
|
||||
|
||||
def test_with_all_sections(self):
|
||||
"""Test with all sections."""
|
||||
result = build_discord_message(
|
||||
base_message="Base",
|
||||
metadata_section="\nMeta",
|
||||
prompt_section="\nPrompt",
|
||||
additional_sections=["\nExtra1", "\nExtra2"]
|
||||
)
|
||||
self.assertIn("Base", result)
|
||||
self.assertIn("Meta", result)
|
||||
self.assertIn("Prompt", result)
|
||||
self.assertIn("Extra1", result)
|
||||
self.assertIn("Extra2", result)
|
||||
|
||||
def test_truncation(self):
|
||||
"""Test message truncation at max length."""
|
||||
long_message = "x" * 2500
|
||||
result = build_discord_message(base_message=long_message, max_length=2000)
|
||||
self.assertLessEqual(len(result), 2000)
|
||||
self.assertIn("[Message truncated]", result)
|
||||
|
||||
def test_no_truncation_under_limit(self):
|
||||
"""Test message not truncated when under limit."""
|
||||
message = "x" * 100
|
||||
result = build_discord_message(base_message=message)
|
||||
self.assertNotIn("truncated", result)
|
||||
|
||||
|
||||
class TestValidateMessageContent(unittest.TestCase):
|
||||
"""Test validate_message_content function."""
|
||||
|
||||
def test_empty_message(self):
|
||||
"""Test empty message is valid."""
|
||||
is_valid, msg = validate_message_content("")
|
||||
self.assertTrue(is_valid)
|
||||
self.assertIn("Empty message", msg)
|
||||
|
||||
def test_normal_message(self):
|
||||
"""Test normal message is valid."""
|
||||
is_valid, msg = validate_message_content("Hello world")
|
||||
self.assertTrue(is_valid)
|
||||
|
||||
def test_too_long_message(self):
|
||||
"""Test message over 2000 chars is invalid."""
|
||||
is_valid, msg = validate_message_content("x" * 2001)
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("2000 character limit", msg)
|
||||
|
||||
def test_message_with_prompts_section(self):
|
||||
"""Test message with Generation Prompts section."""
|
||||
message = "Test\n**Generation Prompts:**\nContent"
|
||||
is_valid, msg = validate_message_content(message)
|
||||
self.assertTrue(is_valid)
|
||||
self.assertNotIn("WARNING", msg)
|
||||
|
||||
def test_message_without_prompts_section(self):
|
||||
"""Test message without Generation Prompts section shows warning."""
|
||||
is_valid, msg = validate_message_content("Test message")
|
||||
self.assertTrue(is_valid)
|
||||
self.assertIn("WARNING", msg)
|
||||
|
||||
|
||||
class TestFormatFileSize(unittest.TestCase):
|
||||
"""Test format_file_size function."""
|
||||
|
||||
def test_bytes(self):
|
||||
"""Test formatting bytes."""
|
||||
self.assertEqual(format_file_size(500), "500 bytes")
|
||||
|
||||
def test_kilobytes(self):
|
||||
"""Test formatting kilobytes."""
|
||||
self.assertEqual(format_file_size(2048), "2.0 KB")
|
||||
|
||||
def test_megabytes(self):
|
||||
"""Test formatting megabytes."""
|
||||
self.assertEqual(format_file_size(5 * 1024 * 1024), "5.0 MB")
|
||||
|
||||
def test_gigabytes(self):
|
||||
"""Test formatting gigabytes."""
|
||||
self.assertEqual(format_file_size(2 * 1024 * 1024 * 1024), "2.00 GB")
|
||||
|
||||
def test_zero(self):
|
||||
"""Test formatting zero bytes."""
|
||||
self.assertEqual(format_file_size(0), "0 bytes")
|
||||
|
||||
|
||||
class TestFormatFileInfo(unittest.TestCase):
|
||||
"""Test format_file_info function."""
|
||||
|
||||
def test_basic_info(self):
|
||||
"""Test basic file info formatting."""
|
||||
result = format_file_info("image.png", 1024)
|
||||
self.assertIn("image.png", result)
|
||||
self.assertIn("1.0 KB", result)
|
||||
|
||||
def test_with_mime_type(self):
|
||||
"""Test file info with MIME type."""
|
||||
result = format_file_info("video.mp4", 1024 * 1024, "video/mp4")
|
||||
self.assertIn("video.mp4", result)
|
||||
self.assertIn("1.0 MB", result)
|
||||
self.assertIn("[video/mp4]", result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,110 @@
|
||||
import unittest
|
||||
import subprocess
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Ensure we get real numpy, not a mock from other test files
|
||||
# Remove any mocked numpy before importing
|
||||
if 'numpy' in sys.modules and hasattr(sys.modules['numpy'], '_mock_name'):
|
||||
del sys.modules['numpy']
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Verify numpy is real
|
||||
assert hasattr(np, 'arange'), "numpy.arange not found - numpy may be mocked"
|
||||
|
||||
class TestNumpyToSubprocess(unittest.TestCase):
|
||||
"""
|
||||
Verify that subprocess.Popen.stdin.write accepts numpy arrays directly.
|
||||
For subprocess.run(input=...), we need to be careful with numpy arrays due to ambiguity check in subprocess module.
|
||||
"""
|
||||
|
||||
def test_popen_stdin_write_numpy(self):
|
||||
"""Test writing numpy array to Popen.stdin"""
|
||||
# Create a small numpy array
|
||||
data = np.arange(256, dtype=np.uint8)
|
||||
|
||||
# Use python to echo input to output (cross-platform)
|
||||
cmd = [sys.executable, '-c', 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())']
|
||||
p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
|
||||
|
||||
# Write numpy array directly
|
||||
p.stdin.write(data)
|
||||
out, _ = p.communicate()
|
||||
|
||||
# Verify output matches input data bytes
|
||||
self.assertEqual(out, data.tobytes())
|
||||
self.assertEqual(len(out), 256)
|
||||
|
||||
def test_run_input_memoryview(self):
|
||||
"""
|
||||
Test passing numpy array as memoryview to subprocess.run input.
|
||||
"""
|
||||
data = np.arange(256, dtype=np.uint8)
|
||||
|
||||
# Use python to echo input to output (cross-platform)
|
||||
cmd = [sys.executable, '-c', 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())']
|
||||
# memoryview works and avoids copy
|
||||
res = subprocess.run(cmd, input=memoryview(data), capture_output=True)
|
||||
|
||||
self.assertEqual(res.stdout, data.tobytes())
|
||||
self.assertEqual(len(res.stdout), 256)
|
||||
|
||||
def test_run_input_fixed_non_contiguous(self):
|
||||
"""
|
||||
Test that using ascontiguousarray makes the non-contiguous array accepted by subprocess.run
|
||||
"""
|
||||
# Create a 2D array and transpose it to make it non-contiguous
|
||||
data = np.zeros((10, 10), dtype=np.uint8)
|
||||
# Fill with some data
|
||||
for i in range(10):
|
||||
for j in range(10):
|
||||
data[i, j] = i + j
|
||||
|
||||
# Transpose creates a non-contiguous view
|
||||
transposed_data = data.T
|
||||
self.assertFalse(transposed_data.flags['C_CONTIGUOUS'])
|
||||
|
||||
# Fix it using ascontiguousarray
|
||||
contiguous_data = np.ascontiguousarray(transposed_data)
|
||||
self.assertTrue(contiguous_data.flags['C_CONTIGUOUS'])
|
||||
|
||||
# Now pass to subprocess
|
||||
mv = memoryview(contiguous_data)
|
||||
cmd = [sys.executable, '-c', 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())']
|
||||
res = subprocess.run(cmd, input=mv, capture_output=True)
|
||||
self.assertEqual(res.stdout, contiguous_data.tobytes())
|
||||
|
||||
def test_popen_stdin_write_fixed_non_contiguous(self):
|
||||
"""
|
||||
Test writing fixed (made contiguous) numpy array to Popen.stdin.
|
||||
"""
|
||||
# Create a 2D array and transpose it to make it non-contiguous
|
||||
data = np.zeros((10, 10), dtype=np.uint8)
|
||||
# Fill with some data
|
||||
for i in range(10):
|
||||
for j in range(10):
|
||||
data[i, j] = i + j
|
||||
|
||||
transposed_data = data.T
|
||||
self.assertFalse(transposed_data.flags['C_CONTIGUOUS'])
|
||||
|
||||
# Fix it
|
||||
contiguous_data = np.ascontiguousarray(transposed_data)
|
||||
self.assertTrue(contiguous_data.flags['C_CONTIGUOUS'])
|
||||
|
||||
# Use python to echo input to output (cross-platform)
|
||||
cmd = [sys.executable, '-c', 'import sys; sys.stdout.buffer.write(sys.stdin.buffer.read())']
|
||||
p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
|
||||
|
||||
try:
|
||||
# Should succeed now
|
||||
p.stdin.write(contiguous_data)
|
||||
out, _ = p.communicate()
|
||||
self.assertEqual(out, contiguous_data.tobytes())
|
||||
|
||||
except Exception as e:
|
||||
self.fail(f"Caught unexpected exception: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,75 @@
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock dependencies to allow testing in isolation
|
||||
sys.modules["folder_paths"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["server"] = MagicMock()
|
||||
sys.modules["requests"] = MagicMock()
|
||||
sys.modules["PIL"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["torch.nn"] = MagicMock()
|
||||
sys.modules["torch.nn.functional"] = MagicMock()
|
||||
|
||||
# Add project root to sys.path
|
||||
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if project_root not in sys.path:
|
||||
sys.path.insert(0, project_root)
|
||||
|
||||
from shared.path_utils import validate_path_is_safe
|
||||
|
||||
class TestPathSecurity(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.test_dir = os.path.join(os.getcwd(), "test_safe_env")
|
||||
os.makedirs(self.test_dir, exist_ok=True)
|
||||
|
||||
def tearDown(self):
|
||||
# Clean up would go here, but since we use temp dirs or mocks, it's fine.
|
||||
# Ideally use tempfile.TemporaryDirectory but this is simple.
|
||||
import shutil
|
||||
if os.path.exists(self.test_dir):
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def test_absolute_path_blocked_with_base_dir(self):
|
||||
# This test checks if validate_path_is_safe blocks writing to /tmp
|
||||
# when base_dir is provided.
|
||||
|
||||
base_dir = self.test_dir
|
||||
test_path = "/tmp/sentinel_test_file.txt"
|
||||
|
||||
try:
|
||||
validate_path_is_safe(test_path, base_dir=base_dir)
|
||||
self.fail("Should have raised ValueError")
|
||||
except ValueError as e:
|
||||
self.assertIn("outside the allowed directory", str(e))
|
||||
|
||||
def test_traversal_blocked_with_base_dir(self):
|
||||
# Create a directory structure
|
||||
base_dir = self.test_dir
|
||||
subdir = os.path.join(base_dir, "subdir")
|
||||
os.makedirs(subdir, exist_ok=True)
|
||||
|
||||
# ../ traversal
|
||||
# This path resolves to outside base_dir
|
||||
test_path = os.path.abspath(os.path.join(subdir, "../../test_safe_escape.txt"))
|
||||
|
||||
try:
|
||||
validate_path_is_safe(test_path, base_dir=base_dir)
|
||||
self.fail("Should have raised ValueError")
|
||||
except ValueError as e:
|
||||
self.assertIn("outside the allowed directory", str(e))
|
||||
|
||||
def test_valid_path_allowed_with_base_dir(self):
|
||||
base_dir = self.test_dir
|
||||
test_path = os.path.join(base_dir, "valid_file.txt")
|
||||
|
||||
# Should not raise
|
||||
validate_path_is_safe(test_path, base_dir=base_dir)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Tests for shared/path_utils.py"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.path_utils import (
|
||||
get_output_directory,
|
||||
ensure_directory_exists,
|
||||
get_unique_filepath,
|
||||
)
|
||||
|
||||
|
||||
class TestGetOutputDirectory(unittest.TestCase):
|
||||
"""Test get_output_directory function."""
|
||||
|
||||
def setUp(self):
|
||||
"""Create temporary directories for testing."""
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
self.output_dir = os.path.join(self.test_dir, "output")
|
||||
self.temp_dir = os.path.join(self.test_dir, "temp")
|
||||
os.makedirs(self.output_dir)
|
||||
os.makedirs(self.temp_dir)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directories."""
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def test_save_output_true(self):
|
||||
"""Test output directory when saving is enabled."""
|
||||
result = get_output_directory(
|
||||
save_output=True,
|
||||
comfy_output_dir=self.output_dir,
|
||||
temp_dir=self.temp_dir,
|
||||
)
|
||||
expected = os.path.join(self.output_dir, "discord_output")
|
||||
self.assertEqual(result, expected)
|
||||
self.assertTrue(os.path.exists(result))
|
||||
|
||||
def test_save_output_false(self):
|
||||
"""Test temp directory when saving is disabled."""
|
||||
result = get_output_directory(
|
||||
save_output=False,
|
||||
comfy_output_dir=self.output_dir,
|
||||
temp_dir=self.temp_dir,
|
||||
)
|
||||
self.assertEqual(result, self.temp_dir)
|
||||
|
||||
def test_custom_subfolder(self):
|
||||
"""Test with custom subfolder name."""
|
||||
result = get_output_directory(
|
||||
save_output=True,
|
||||
comfy_output_dir=self.output_dir,
|
||||
temp_dir=self.temp_dir,
|
||||
subfolder="custom_folder",
|
||||
)
|
||||
expected = os.path.join(self.output_dir, "custom_folder")
|
||||
self.assertEqual(result, expected)
|
||||
self.assertTrue(os.path.exists(result))
|
||||
|
||||
|
||||
class TestEnsureDirectoryExists(unittest.TestCase):
|
||||
"""Test ensure_directory_exists function."""
|
||||
|
||||
def setUp(self):
|
||||
"""Create temporary directory for testing."""
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directories."""
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def test_creates_directory(self):
|
||||
"""Test that directory is created if it doesn't exist."""
|
||||
new_dir = os.path.join(self.test_dir, "new_directory")
|
||||
self.assertFalse(os.path.exists(new_dir))
|
||||
result = ensure_directory_exists(new_dir)
|
||||
self.assertTrue(os.path.exists(new_dir))
|
||||
self.assertEqual(result, new_dir)
|
||||
|
||||
def test_existing_directory(self):
|
||||
"""Test that existing directory is not affected."""
|
||||
result = ensure_directory_exists(self.test_dir)
|
||||
self.assertTrue(os.path.exists(self.test_dir))
|
||||
self.assertEqual(result, self.test_dir)
|
||||
|
||||
def test_nested_directories(self):
|
||||
"""Test creating nested directories."""
|
||||
nested = os.path.join(self.test_dir, "a", "b", "c")
|
||||
result = ensure_directory_exists(nested)
|
||||
self.assertTrue(os.path.exists(nested))
|
||||
self.assertEqual(result, nested)
|
||||
|
||||
|
||||
class TestGetUniqueFilepath(unittest.TestCase):
|
||||
"""Test get_unique_filepath function."""
|
||||
|
||||
def test_basic_filepath(self):
|
||||
"""Test basic filepath generation."""
|
||||
result = get_unique_filepath("/output", "image", ".png")
|
||||
self.assertEqual(result, "/output/image.png")
|
||||
|
||||
def test_with_counter(self):
|
||||
"""Test filepath with counter."""
|
||||
result = get_unique_filepath("/output", "image", ".png", counter=5)
|
||||
self.assertEqual(result, "/output/image_00005.png")
|
||||
|
||||
def test_counter_formatting(self):
|
||||
"""Test counter is formatted with leading zeros."""
|
||||
result = get_unique_filepath("/output", "image", ".jpg", counter=123)
|
||||
self.assertEqual(result, "/output/image_00123.jpg")
|
||||
|
||||
def test_extension_without_dot(self):
|
||||
"""Test extension is normalized if dot is missing."""
|
||||
result = get_unique_filepath("/output", "video", "mp4")
|
||||
self.assertEqual(result, "/output/video.mp4")
|
||||
|
||||
def test_extension_with_dot(self):
|
||||
"""Test extension with dot works correctly."""
|
||||
result = get_unique_filepath("/output", "video", ".mp4")
|
||||
self.assertEqual(result, "/output/video.mp4")
|
||||
|
||||
def test_counter_zero(self):
|
||||
"""Test counter value of zero."""
|
||||
result = get_unique_filepath("/output", "frame", ".png", counter=0)
|
||||
self.assertEqual(result, "/output/frame_00000.png")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,151 @@
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
import sys
|
||||
import os
|
||||
import numpy as np
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# Mock dependencies before importing nodes.video_node
|
||||
mock_torch = MagicMock()
|
||||
sys.modules["torch"] = mock_torch
|
||||
sys.modules["folder_paths"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["comfy.cli_args"] = MagicMock()
|
||||
sys.modules["comfy.utils"] = MagicMock()
|
||||
sys.modules["server"] = MagicMock()
|
||||
|
||||
# Mock PIL
|
||||
mock_pil = MagicMock()
|
||||
sys.modules["PIL"] = mock_pil
|
||||
sys.modules["PIL.Image"] = mock_pil
|
||||
sys.modules["PIL.PngImagePlugin"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Mock shared modules' submodules to avoid dependency issues
|
||||
sys.modules["shared.workflow"] = MagicMock()
|
||||
sys.modules["shared.workflow.sanitizer"] = MagicMock()
|
||||
sys.modules["shared.workflow.prompt_extractor"] = MagicMock()
|
||||
sys.modules["shared.workflow.workflow_builder"] = MagicMock()
|
||||
|
||||
sys.modules["shared.discord"] = MagicMock()
|
||||
sys.modules["shared.discord.webhook_client"] = MagicMock()
|
||||
sys.modules["shared.discord.message_builder"] = MagicMock()
|
||||
sys.modules["shared.discord.cdn_extractor"] = MagicMock()
|
||||
|
||||
sys.modules["shared.github_integration"] = MagicMock()
|
||||
sys.modules["shared.logging_config"] = MagicMock()
|
||||
sys.modules["shared.filename_utils"] = MagicMock()
|
||||
sys.modules["shared.path_utils"] = MagicMock()
|
||||
|
||||
# Mock shared.media siblings
|
||||
sys.modules["shared.media.format_utils"] = MagicMock()
|
||||
sys.modules["shared.media.video_encoder"] = MagicMock()
|
||||
|
||||
# Note: We do NOT mock "shared", "shared.media", or "shared.media.image_processing"
|
||||
# because we want to load the real code for testing.
|
||||
|
||||
# Define the function logic we want to verify (simulating the generator consumer)
|
||||
def consume_chunks(chunks):
|
||||
pil_images = []
|
||||
for chunk in chunks:
|
||||
# Proposed logic for nodes/video_node.py
|
||||
if len(chunk.shape) == 4:
|
||||
# Batched chunk (B, H, W, C)
|
||||
for i in range(chunk.shape[0]):
|
||||
pil_images.append(f"image_from_batch_{i}")
|
||||
else:
|
||||
# Single frame chunk (H, W, C)
|
||||
pil_images.append("image_from_single")
|
||||
return pil_images
|
||||
|
||||
class TestPILBatchOptimization(unittest.TestCase):
|
||||
|
||||
def test_consumer_logic_mixed_chunks(self):
|
||||
"""Test that the consumer logic correctly handles mixed 4D and 3D chunks."""
|
||||
|
||||
# 1. 4D Chunk (Batch of 2)
|
||||
chunk_batch = np.zeros((2, 10, 10, 3), dtype=np.uint8)
|
||||
|
||||
# 2. 3D Chunk (Single frame) - simulating what happens if generator yields single frame
|
||||
chunk_single = np.zeros((10, 10, 3), dtype=np.uint8)
|
||||
|
||||
chunks = [chunk_batch, chunk_single]
|
||||
|
||||
# Run consumer logic
|
||||
images = consume_chunks(chunks)
|
||||
|
||||
# Verify results
|
||||
# Should have 2 from batch + 1 from single = 3 images
|
||||
self.assertEqual(len(images), 3)
|
||||
self.assertEqual(images[0], "image_from_batch_0")
|
||||
self.assertEqual(images[1], "image_from_batch_1")
|
||||
self.assertEqual(images[2], "image_from_single")
|
||||
|
||||
def test_process_batched_images_integration(self):
|
||||
"""
|
||||
Verify that we can import and run process_batched_images with mocks,
|
||||
and that it chunks correctly.
|
||||
"""
|
||||
# Import needs to happen after mocks are set up
|
||||
from shared.media.image_processing import process_batched_images
|
||||
|
||||
# Setup mock tensor
|
||||
# We need to make sure isinstance(t, torch.Tensor) works
|
||||
|
||||
tensor_len = 5
|
||||
batch_size = 2
|
||||
|
||||
# Mock slicing
|
||||
def getitem(self, idx):
|
||||
# idx is a slice object
|
||||
start = idx.start
|
||||
stop = idx.stop
|
||||
if stop > tensor_len:
|
||||
stop = tensor_len
|
||||
size = stop - start
|
||||
return f"slice_{size}"
|
||||
|
||||
# Create a class with __len__ and __getitem__ defined
|
||||
class MockTensor:
|
||||
def __len__(self):
|
||||
return tensor_len
|
||||
def __getitem__(self, idx):
|
||||
return getitem(self, idx)
|
||||
|
||||
mock_torch.Tensor = MockTensor
|
||||
|
||||
mock_tensor = mock_torch.Tensor()
|
||||
|
||||
# Mock tensor_to_numpy_uint8 to return numpy arrays of appropriate shape
|
||||
# It needs to return (Size, H, W, C)
|
||||
with patch('shared.media.image_processing.tensor_to_numpy_uint8') as mock_t2n:
|
||||
def side_effect(slice_obj):
|
||||
# parse size from string "slice_N"
|
||||
size = int(slice_obj.split('_')[1])
|
||||
return np.zeros((size, 10, 10, 3), dtype=np.uint8)
|
||||
|
||||
mock_t2n.side_effect = side_effect
|
||||
|
||||
# Run generator
|
||||
generator = process_batched_images(mock_tensor, batch_size=batch_size)
|
||||
chunks = list(generator)
|
||||
|
||||
# Expected:
|
||||
# 5 items, batch 2
|
||||
# 1. Size 2
|
||||
# 2. Size 2
|
||||
# 3. Size 1
|
||||
|
||||
self.assertEqual(len(chunks), 3)
|
||||
self.assertEqual(chunks[0].shape[0], 2)
|
||||
self.assertEqual(chunks[1].shape[0], 2)
|
||||
self.assertEqual(chunks[2].shape[0], 1)
|
||||
|
||||
# Verify they are all 4D arrays (B, H, W, C)
|
||||
for c in chunks:
|
||||
self.assertEqual(len(c.shape), 4)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,125 @@
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
sys.modules["PIL"] = MagicMock()
|
||||
sys.modules["PIL.Image"] = MagicMock()
|
||||
sys.modules["PIL.PngImagePlugin"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["comfy.cli_args"] = MagicMock()
|
||||
sys.modules["comfy.utils"] = MagicMock()
|
||||
sys.modules["server"] = MagicMock()
|
||||
|
||||
# Mock folder_paths
|
||||
mock_folder_paths = MagicMock()
|
||||
sys.modules["folder_paths"] = mock_folder_paths
|
||||
|
||||
# Add parent directory to path
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
# Import the node
|
||||
from nodes.video_node import DiscordSendSaveVideo
|
||||
|
||||
class TestTempFileLeak(unittest.TestCase):
|
||||
def setUp(self):
|
||||
# Create a real temporary directory for our test
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
self.output_dir = os.path.join(self.test_dir, "output")
|
||||
self.temp_dir = os.path.join(self.test_dir, "temp")
|
||||
os.makedirs(self.output_dir)
|
||||
os.makedirs(self.temp_dir)
|
||||
|
||||
# Configure folder_paths mock
|
||||
mock_folder_paths.get_output_directory.return_value = self.output_dir
|
||||
mock_folder_paths.get_temp_directory.return_value = self.temp_dir
|
||||
|
||||
# Mock get_save_image_path to return predictable paths
|
||||
# full_output_folder, filename, counter, subfolder, filename_prefix
|
||||
mock_folder_paths.get_save_image_path.return_value = (
|
||||
self.output_dir, "ComfyUI-Video", 1, "", "ComfyUI-Video"
|
||||
)
|
||||
|
||||
# Instantiate the node
|
||||
self.node = DiscordSendSaveVideo()
|
||||
|
||||
# Create a dummy image tensor mock
|
||||
self.dummy_image = MagicMock()
|
||||
self.dummy_image.shape = (512, 512, 3) # height, width, channels
|
||||
|
||||
# Mock tensor_to_numpy_uint8 in discordsend_utils
|
||||
self.patcher_numpy = patch("nodes.video_node.tensor_to_numpy_uint8")
|
||||
self.mock_numpy_conv = self.patcher_numpy.start()
|
||||
# Return a dummy numpy array
|
||||
import numpy as np
|
||||
self.mock_numpy_conv.return_value = np.zeros((512, 512, 3), dtype=np.uint8)
|
||||
|
||||
def tearDown(self):
|
||||
self.patcher_numpy.stop()
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
@patch("nodes.video_node.subprocess.Popen")
|
||||
@patch("nodes.video_node.subprocess.run")
|
||||
@patch("nodes.base_node.send_to_discord_with_retry")
|
||||
@patch("nodes.video_node.Image")
|
||||
@patch("nodes.video_node.os.path.getsize")
|
||||
@patch("nodes.video_node.validate_video_for_discord")
|
||||
def test_temp_file_leak(self, mock_validate, mock_getsize, mock_image, mock_send, mock_run, mock_popen):
|
||||
# Setup mocks
|
||||
mock_process = MagicMock()
|
||||
mock_process.returncode = 0
|
||||
mock_process.stdin = MagicMock()
|
||||
mock_popen.return_value = mock_process
|
||||
|
||||
mock_send.return_value.status_code = 200
|
||||
mock_send.return_value.json.return_value = {}
|
||||
|
||||
mock_getsize.return_value = 1024 * 1024 # 1MB
|
||||
mock_validate.return_value = (True, "Valid")
|
||||
|
||||
# Create a fake output file that subprocess would have created
|
||||
fake_output_path = os.path.join(self.output_dir, "ComfyUI-Video_00001.mp4")
|
||||
with open(fake_output_path, "wb") as f:
|
||||
f.write(b"fake video content")
|
||||
|
||||
# Mock subprocess.run to simulate creation of optimized file
|
||||
def side_effect_run(args, **kwargs):
|
||||
# The last argument is the output file path
|
||||
output_file = args[-1]
|
||||
if "discord_optimized_" in output_file:
|
||||
# Create the file
|
||||
with open(output_file, "wb") as f:
|
||||
f.write(b"optimized video content")
|
||||
return MagicMock(returncode=0)
|
||||
|
||||
mock_run.side_effect = side_effect_run
|
||||
|
||||
# Run the node
|
||||
self.node.save_video(
|
||||
images=[self.dummy_image],
|
||||
send_to_discord=True,
|
||||
webhook_url="https://discord.com/api/webhooks/123/abc",
|
||||
format="video/h264-mp4",
|
||||
save_output=True
|
||||
)
|
||||
|
||||
# Check if any file in temp_dir contains "discord_optimized_"
|
||||
temp_files = os.listdir(self.temp_dir)
|
||||
optimized_files = [f for f in temp_files if "discord_optimized_" in f]
|
||||
|
||||
print(f"Files remaining in temp dir: {temp_files}")
|
||||
|
||||
# This assertion is expected to FAIL if the leak exists (because we want 0 files)
|
||||
# Or pass if we assert > 0 to prove the leak.
|
||||
# After fix, we expect 0 files.
|
||||
self.assertEqual(len(optimized_files), 0, "Temporary optimized file should have been cleaned up")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,246 @@
|
||||
import unittest
|
||||
import sys
|
||||
import os
|
||||
import tempfile
|
||||
import shutil
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Add project root to sys.path
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
# Mock comfy modules
|
||||
if 'comfy' not in sys.modules:
|
||||
sys.modules['comfy'] = MagicMock()
|
||||
if 'comfy.cli_args' not in sys.modules:
|
||||
sys.modules['comfy.cli_args'] = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args.disable_metadata = False
|
||||
if 'comfy.utils' not in sys.modules:
|
||||
sys.modules['comfy.utils'] = MagicMock()
|
||||
if 'folder_paths' not in sys.modules:
|
||||
sys.modules['folder_paths'] = MagicMock()
|
||||
if 'server' not in sys.modules:
|
||||
sys.modules['server'] = MagicMock()
|
||||
|
||||
# Mock heavy/external dependencies
|
||||
sys.modules['torch'] = MagicMock()
|
||||
sys.modules['cv2'] = MagicMock()
|
||||
# We rely on real Pillow and numpy being installed and used
|
||||
|
||||
# Import the node
|
||||
try:
|
||||
from nodes.video_node import DiscordSendSaveVideo
|
||||
except ImportError:
|
||||
raise
|
||||
|
||||
class TestSymlinkAttack(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.test_dir = tempfile.mkdtemp()
|
||||
self.output_dir = os.path.join(self.test_dir, "output")
|
||||
self.temp_dir = os.path.join(self.test_dir, "temp")
|
||||
os.makedirs(self.output_dir)
|
||||
os.makedirs(self.temp_dir)
|
||||
|
||||
self.node = DiscordSendSaveVideo()
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.test_dir)
|
||||
|
||||
def test_overwrite_symlink_vulnerability(self):
|
||||
"""Test overwriting a direct symlink file."""
|
||||
# Create a target file
|
||||
target_file = os.path.join(self.test_dir, "target.txt")
|
||||
with open(target_file, "w") as f:
|
||||
f.write("Original Content")
|
||||
|
||||
# Create symlinks in output dir pointing to target file
|
||||
symlink_video = os.path.join(self.output_dir, "ComfyUI-Video_00001.mp4")
|
||||
os.symlink(target_file, symlink_video)
|
||||
|
||||
# Patch folder_paths on the node module
|
||||
with patch('nodes.video_node.folder_paths') as mock_folder_paths:
|
||||
mock_folder_paths.get_save_image_path.return_value = (
|
||||
self.output_dir,
|
||||
"ComfyUI-Video",
|
||||
1,
|
||||
"",
|
||||
"ComfyUI-Video"
|
||||
)
|
||||
|
||||
# Mock images input
|
||||
mock_images = [MagicMock()]
|
||||
mock_images[0].shape = (64, 64, 3)
|
||||
|
||||
with patch('nodes.video_node.tensor_to_numpy_uint8') as mock_t2n:
|
||||
import numpy as np
|
||||
mock_t2n.return_value = np.zeros((64, 64, 3), dtype=np.uint8)
|
||||
|
||||
with patch('nodes.video_node.ffmpeg_path', "ffmpeg"):
|
||||
with patch('subprocess.Popen') as mock_popen:
|
||||
mock_popen.return_value = MagicMock()
|
||||
|
||||
# Expect ValueError when trying to overwrite symlink
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.node.save_video(
|
||||
images=mock_images,
|
||||
overwrite_last=True,
|
||||
format="video/h264-mp4",
|
||||
save_output=True,
|
||||
frame_rate=1.0
|
||||
)
|
||||
|
||||
self.assertIn("symlink", str(context.exception))
|
||||
print("SUCCESS: Direct symlink overwrite prevented.")
|
||||
|
||||
def test_parent_directory_symlink_vulnerability(self):
|
||||
"""Test writing to a path where a parent directory is a symlink."""
|
||||
# Create a real directory outside the intended output
|
||||
secret_dir = os.path.join(self.test_dir, "secret")
|
||||
os.makedirs(secret_dir)
|
||||
|
||||
# Create a symlink inside output_dir pointing to secret_dir
|
||||
# /test_dir/output/evil_link -> /test_dir/secret
|
||||
evil_link = os.path.join(self.output_dir, "evil_link")
|
||||
os.symlink(secret_dir, evil_link)
|
||||
|
||||
# We want to write to /test_dir/output/evil_link/file.mp4
|
||||
# which resolves to /test_dir/secret/file.mp4
|
||||
|
||||
# Patch folder_paths to return the evil_link as the directory
|
||||
with patch('nodes.video_node.folder_paths') as mock_folder_paths:
|
||||
mock_folder_paths.get_save_image_path.return_value = (
|
||||
evil_link,
|
||||
"ComfyUI-Video",
|
||||
1,
|
||||
"",
|
||||
"ComfyUI-Video"
|
||||
)
|
||||
|
||||
mock_images = [MagicMock()]
|
||||
mock_images[0].shape = (64, 64, 3)
|
||||
|
||||
with patch('nodes.video_node.tensor_to_numpy_uint8') as mock_t2n:
|
||||
import numpy as np
|
||||
mock_t2n.return_value = np.zeros((64, 64, 3), dtype=np.uint8)
|
||||
|
||||
with patch('nodes.video_node.ffmpeg_path', "ffmpeg"):
|
||||
with patch('subprocess.Popen') as mock_popen:
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.node.save_video(
|
||||
images=mock_images,
|
||||
overwrite_last=True,
|
||||
format="video/h264-mp4",
|
||||
save_output=True,
|
||||
frame_rate=1.0
|
||||
)
|
||||
|
||||
# Check for either the parent dir message or the generic mismatch message
|
||||
error_msg = str(context.exception)
|
||||
self.assertTrue(
|
||||
"Writing through directory symlinks is not allowed" in error_msg or
|
||||
"Symlinks in output paths are not allowed" in error_msg,
|
||||
f"Unexpected error message: {error_msg}"
|
||||
)
|
||||
print("SUCCESS: Parent directory symlink prevented.")
|
||||
|
||||
def test_non_existent_directory_symlink_bypass(self):
|
||||
"""Test where intermediate directory doesn't exist but parent is symlink."""
|
||||
# /test_dir/secret
|
||||
secret_dir = os.path.join(self.test_dir, "secret")
|
||||
os.makedirs(secret_dir)
|
||||
|
||||
# /test_dir/output/link -> /test_dir/secret
|
||||
link_dir = os.path.join(self.output_dir, "link")
|
||||
os.symlink(secret_dir, link_dir)
|
||||
|
||||
# Target: /test_dir/output/link/subdir/file.mp4
|
||||
# 'subdir' does not exist yet.
|
||||
target_dir = os.path.join(link_dir, "subdir")
|
||||
# Do NOT create target_dir.
|
||||
|
||||
# Patch folder_paths to return the non-existent target_dir
|
||||
with patch('nodes.video_node.folder_paths') as mock_folder_paths:
|
||||
mock_folder_paths.get_save_image_path.return_value = (
|
||||
target_dir,
|
||||
"ComfyUI-Video",
|
||||
1,
|
||||
"",
|
||||
"ComfyUI-Video"
|
||||
)
|
||||
|
||||
# Use os.makedirs real implementation to create the directory if the node calls it
|
||||
# But here we assume the validation happens before directory creation or during path validation
|
||||
|
||||
mock_images = [MagicMock()]
|
||||
mock_images[0].shape = (64, 64, 3)
|
||||
|
||||
with patch('nodes.video_node.tensor_to_numpy_uint8') as mock_t2n:
|
||||
import numpy as np
|
||||
mock_t2n.return_value = np.zeros((64, 64, 3), dtype=np.uint8)
|
||||
|
||||
with patch('nodes.video_node.ffmpeg_path', "ffmpeg"):
|
||||
with patch('subprocess.Popen') as mock_popen:
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.node.save_video(
|
||||
images=mock_images,
|
||||
overwrite_last=True,
|
||||
format="video/h264-mp4",
|
||||
save_output=True,
|
||||
frame_rate=1.0
|
||||
)
|
||||
|
||||
error_msg = str(context.exception)
|
||||
self.assertTrue(
|
||||
"Symlinks in output paths are not allowed" in error_msg or
|
||||
"Path component" in error_msg and "is a symlink" in error_msg,
|
||||
f"Unexpected error message: {error_msg}"
|
||||
)
|
||||
print("SUCCESS: Non-existent directory symlink bypass prevented.")
|
||||
|
||||
def test_vhs_format_bypass(self):
|
||||
"""Test that VHS format path recalculation is also validated."""
|
||||
# This test tries to exploit the path where 'is_vhs_format' is True
|
||||
# which changes the file extension and potentially bypasses early checks
|
||||
|
||||
target_file = os.path.join(self.test_dir, "target.mkv")
|
||||
with open(target_file, "w") as f:
|
||||
f.write("Original Content")
|
||||
|
||||
# Create symlink with different extension (mkv) that VHS might use
|
||||
symlink_video = os.path.join(self.output_dir, "ComfyUI-Video_00001.mkv")
|
||||
os.symlink(target_file, symlink_video)
|
||||
|
||||
with patch('nodes.video_node.folder_paths') as mock_folder_paths:
|
||||
mock_folder_paths.get_save_image_path.return_value = (
|
||||
self.output_dir,
|
||||
"ComfyUI-Video",
|
||||
1,
|
||||
"",
|
||||
"ComfyUI-Video"
|
||||
)
|
||||
|
||||
mock_images = [MagicMock()]
|
||||
mock_images[0].shape = (64, 64, 3)
|
||||
|
||||
with patch('nodes.video_node.tensor_to_numpy_uint8') as mock_t2n:
|
||||
import numpy as np
|
||||
mock_t2n.return_value = np.zeros((64, 64, 3), dtype=np.uint8)
|
||||
|
||||
with patch('nodes.video_node.ffmpeg_path', "ffmpeg"):
|
||||
# Mock has_vhs_formats to be True
|
||||
with patch('nodes.video_node.has_vhs_formats', True):
|
||||
with patch('subprocess.Popen') as mock_popen:
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.node.save_video(
|
||||
images=mock_images,
|
||||
overwrite_last=True,
|
||||
format="video/mkv", # Custom format triggers VHS path
|
||||
save_output=True,
|
||||
frame_rate=1.0
|
||||
)
|
||||
|
||||
self.assertIn("symlink", str(context.exception))
|
||||
print("SUCCESS: VHS format path bypass prevented.")
|
||||
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
Unit tests for ComfyUI-DiscordSend utilities.
|
||||
|
||||
Run tests with: python -m pytest tests/test_utils.py -v
|
||||
Or without pytest: python tests/test_utils.py
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
import requests
|
||||
|
||||
# Mock dependencies before importing project modules
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
# sys.modules["PIL"] = MagicMock() # PIL might be installed, so maybe not mock it if not needed, but safer to mock if we don't rely on it for these tests
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.workflow.sanitizer import sanitize_json_for_export
|
||||
from shared.discord.webhook_client import validate_webhook_url, sanitize_webhook_for_logging, send_to_discord_with_retry, DiscordWebhookClient
|
||||
from shared.github_integration import update_github_cdn_urls
|
||||
|
||||
|
||||
class TestSanitizer(unittest.TestCase):
|
||||
"""Tests for the sanitize_json_for_export function."""
|
||||
|
||||
def test_sanitize_none(self):
|
||||
"""Should handle None input."""
|
||||
self.assertIsNone(sanitize_json_for_export(None))
|
||||
|
||||
def test_sanitize_webhook_url_in_string(self):
|
||||
"""Should remove webhook URLs from strings."""
|
||||
test_url = "https://discord.com/api/webhooks/123456789/abcdefg"
|
||||
result = sanitize_json_for_export(test_url)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_sanitize_github_token_in_string(self):
|
||||
"""Should remove GitHub tokens from strings."""
|
||||
test_token = "ghp_1234567890abcdefghijklmnopqrstuvwxyz"
|
||||
result = sanitize_json_for_export(test_token)
|
||||
self.assertEqual(result, "")
|
||||
|
||||
def test_sanitize_dict_with_webhook_key(self):
|
||||
"""Should remove values for webhook-related keys."""
|
||||
test_data = {
|
||||
"webhook_url": "https://discord.com/api/webhooks/123/abc",
|
||||
"message": "Hello world"
|
||||
}
|
||||
result = sanitize_json_for_export(test_data)
|
||||
self.assertEqual(result["webhook_url"], "")
|
||||
self.assertEqual(result["message"], "Hello world")
|
||||
|
||||
def test_sanitize_dict_with_github_token_key(self):
|
||||
"""Should remove values for github_token keys."""
|
||||
test_data = {
|
||||
"github_token": "ghp_secret123",
|
||||
"repo": "user/repo"
|
||||
}
|
||||
result = sanitize_json_for_export(test_data)
|
||||
self.assertEqual(result["github_token"], "")
|
||||
self.assertEqual(result["repo"], "user/repo")
|
||||
|
||||
def test_sanitize_nested_dict(self):
|
||||
"""Should sanitize nested dictionaries."""
|
||||
test_data = {
|
||||
"level1": {
|
||||
"webhook_url": "https://discord.com/api/webhooks/123/abc",
|
||||
"level2": {
|
||||
"github_token": "ghp_secret"
|
||||
}
|
||||
}
|
||||
}
|
||||
result = sanitize_json_for_export(test_data)
|
||||
self.assertEqual(result["level1"]["webhook_url"], "")
|
||||
self.assertEqual(result["level1"]["level2"]["github_token"], "")
|
||||
|
||||
def test_sanitize_list(self):
|
||||
"""Should sanitize lists."""
|
||||
test_data = [
|
||||
"normal string",
|
||||
"https://discord.com/api/webhooks/123/abc",
|
||||
{"webhook_url": "secret"}
|
||||
]
|
||||
result = sanitize_json_for_export(test_data)
|
||||
self.assertEqual(result[0], "normal string")
|
||||
self.assertEqual(result[1], "")
|
||||
self.assertEqual(result[2]["webhook_url"], "")
|
||||
|
||||
def test_preserve_normal_data(self):
|
||||
"""Should preserve non-sensitive data."""
|
||||
test_data = {
|
||||
"prompt": "a beautiful sunset",
|
||||
"seed": 12345,
|
||||
"steps": 20,
|
||||
"cfg": 7.5
|
||||
}
|
||||
result = sanitize_json_for_export(test_data)
|
||||
self.assertEqual(result, test_data)
|
||||
|
||||
def test_sanitize_node_other_properties(self):
|
||||
"""Should sanitize other properties in nodes that are not inputs/widgets."""
|
||||
test_data = {
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "SomeNode",
|
||||
"inputs": {},
|
||||
"widgets_values": [],
|
||||
"extra": {
|
||||
"webhook_url": "https://discord.com/api/webhooks/123/abc"
|
||||
},
|
||||
"properties": {
|
||||
"nested": {
|
||||
"token": "ghp_secret"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
result = sanitize_json_for_export(test_data)
|
||||
node = result["nodes"][0]
|
||||
self.assertEqual(node["extra"]["webhook_url"], "")
|
||||
self.assertEqual(node["properties"]["nested"]["token"], "")
|
||||
|
||||
|
||||
class TestWebhookValidation(unittest.TestCase):
|
||||
"""Tests for webhook URL validation."""
|
||||
|
||||
def test_valid_webhook_url(self):
|
||||
"""Should accept valid Discord webhook URLs."""
|
||||
valid_url = "https://discord.com/api/webhooks/1234567890/abcdefghijklmnop"
|
||||
is_valid, message = validate_webhook_url(valid_url)
|
||||
self.assertTrue(is_valid)
|
||||
|
||||
def test_valid_discordapp_url(self):
|
||||
"""Should accept discordapp.com URLs."""
|
||||
valid_url = "https://discordapp.com/api/webhooks/1234567890/abcdefghijklmnop"
|
||||
is_valid, message = validate_webhook_url(valid_url)
|
||||
self.assertTrue(is_valid)
|
||||
|
||||
def test_empty_url(self):
|
||||
"""Should reject empty URLs."""
|
||||
is_valid, message = validate_webhook_url("")
|
||||
self.assertFalse(is_valid)
|
||||
|
||||
def test_invalid_url(self):
|
||||
"""Should reject non-webhook URLs."""
|
||||
is_valid, message = validate_webhook_url("https://example.com")
|
||||
self.assertFalse(is_valid)
|
||||
|
||||
def test_bypass_attempt_url(self):
|
||||
"""Should reject URLs that attempt to bypass validation."""
|
||||
# This URL contains 'discord' and 'webhook' but is not hosted on discord.com
|
||||
bypass_url = "http://evil-site.com/discord/webhook"
|
||||
is_valid, message = validate_webhook_url(bypass_url)
|
||||
self.assertFalse(is_valid)
|
||||
|
||||
def test_localhost_url(self):
|
||||
"""Should reject localhost URLs (SSRF protection)."""
|
||||
is_valid, message = validate_webhook_url("http://localhost:8080/admin")
|
||||
self.assertFalse(is_valid)
|
||||
|
||||
def test_http_url_rejected(self):
|
||||
"""Should reject HTTP URLs (must be HTTPS)."""
|
||||
is_valid, message = validate_webhook_url("http://discord.com/api/webhooks/123/abc")
|
||||
self.assertFalse(is_valid)
|
||||
self.assertIn("must start with https://", message)
|
||||
|
||||
def test_ip_encoding_urls(self):
|
||||
"""Should reject alternate IP encodings."""
|
||||
self.assertFalse(validate_webhook_url("http://127.0.0.1")[0])
|
||||
self.assertFalse(validate_webhook_url("http://0177.0.0.1")[0]) # Octal
|
||||
self.assertFalse(validate_webhook_url("http://0x7f.0.0.1")[0]) # Hex
|
||||
self.assertFalse(validate_webhook_url("http://[::1]")[0]) # IPv6
|
||||
|
||||
def test_domain_spoofing_urls(self):
|
||||
"""Should reject domains that contain 'discord' but aren't Discord."""
|
||||
self.assertFalse(validate_webhook_url("https://discord.com.evil.co/api/webhooks/123/abc")[0])
|
||||
self.assertFalse(validate_webhook_url("https://evil-discord.com/api/webhooks/123/abc")[0])
|
||||
|
||||
def test_path_traversal_urls(self):
|
||||
"""Should reject path traversal attempts."""
|
||||
self.assertFalse(validate_webhook_url("https://discord.com/api/webhooks/123/abc/../../admin")[0])
|
||||
self.assertFalse(validate_webhook_url("https://discord.com/api/webhooks/123/abc%2f..%2f..%2fadmin")[0])
|
||||
|
||||
|
||||
class TestSSRFPrevention(unittest.TestCase):
|
||||
"""Tests for SSRF prevention mechanisms."""
|
||||
|
||||
def test_send_to_discord_validates_url(self):
|
||||
"""Should raise ValueError for invalid URLs before sending request."""
|
||||
malicious_url = "http://localhost:8080/admin/delete"
|
||||
|
||||
# We don't need to mock requests.post because it should fail before calling it
|
||||
with self.assertRaises(ValueError) as cm:
|
||||
send_to_discord_with_retry(malicious_url, data={"content": "test"})
|
||||
|
||||
self.assertIn("Invalid webhook URL", str(cm.exception))
|
||||
|
||||
@patch('shared.discord.webhook_client.requests.post')
|
||||
def test_send_to_discord_allows_valid_url(self, mock_post):
|
||||
"""Should allow valid Discord URLs."""
|
||||
valid_url = "https://discord.com/api/webhooks/123/abc"
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
send_to_discord_with_retry(valid_url, data={"content": "test"})
|
||||
|
||||
mock_post.assert_called_once()
|
||||
|
||||
|
||||
class TestWebhookSanitization(unittest.TestCase):
|
||||
"""Tests for webhook URL sanitization for logging."""
|
||||
|
||||
def test_sanitize_webhook_for_logging(self):
|
||||
"""Should redact the token portion of webhook URLs."""
|
||||
url = "https://discord.com/api/webhooks/123456789/secrettoken"
|
||||
result = sanitize_webhook_for_logging(url)
|
||||
self.assertIn("[REDACTED]", result)
|
||||
self.assertIn("123456789", result)
|
||||
self.assertNotIn("secrettoken", result)
|
||||
|
||||
def test_sanitize_empty_url(self):
|
||||
"""Should handle empty URLs."""
|
||||
result = sanitize_webhook_for_logging("")
|
||||
self.assertEqual(result, "")
|
||||
|
||||
|
||||
class TestDiscordWebhookClient(unittest.TestCase):
|
||||
"""Tests for DiscordWebhookClient security features."""
|
||||
|
||||
@patch('shared.discord.webhook_client.requests.post')
|
||||
def test_exception_token_leakage(self, mock_post):
|
||||
"""Should redact tokens from exception messages in last_error."""
|
||||
token = "SUPER_SECRET_TOKEN"
|
||||
url = f"https://discord.com/api/webhooks/123456/{token}"
|
||||
client = DiscordWebhookClient(url)
|
||||
|
||||
# Configure mock to raise an exception containing the token
|
||||
error_message = f"Max retries exceeded with url: /api/webhooks/123456/{token}"
|
||||
mock_post.side_effect = requests.exceptions.ConnectionError(error_message)
|
||||
|
||||
success, result = client.send_message("Test message")
|
||||
|
||||
self.assertFalse(success)
|
||||
self.assertIn("error", result)
|
||||
self.assertNotIn(token, result["error"])
|
||||
self.assertIn("[REDACTED]", result["error"])
|
||||
|
||||
|
||||
class TestGitHubIntegration(unittest.TestCase):
|
||||
"""Tests for GitHub integration security features."""
|
||||
|
||||
@patch('shared.github_integration.requests.put')
|
||||
@patch('shared.github_integration.requests.get')
|
||||
def test_github_token_redaction_in_response(self, mock_get, mock_put):
|
||||
"""Should redact GitHub token from error messages including response text."""
|
||||
token = "ghp_SECRET_TOKEN"
|
||||
repo = "user/repo"
|
||||
file_path = "cdn_urls.md"
|
||||
|
||||
# Mock GET to return 404 (file doesn't exist)
|
||||
mock_get_response = MagicMock()
|
||||
mock_get_response.status_code = 404
|
||||
mock_get.return_value = mock_get_response
|
||||
|
||||
# Mock PUT to fail and return the token in the response text (simulating leak)
|
||||
mock_put_response = MagicMock()
|
||||
mock_put_response.status_code = 401
|
||||
mock_put_response.text = f"Bad credentials: {token} is invalid"
|
||||
mock_put.return_value = mock_put_response
|
||||
|
||||
success, message = update_github_cdn_urls(repo, token, file_path, [("test.png", "http://url")])
|
||||
|
||||
self.assertFalse(success)
|
||||
self.assertNotIn(token, message)
|
||||
self.assertIn("[REDACTED_TOKEN]", message)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run tests
|
||||
print("Running ComfyUI-DiscordSend utility tests...\n")
|
||||
unittest.main(verbosity=2)
|
||||
@@ -0,0 +1,106 @@
|
||||
import unittest
|
||||
import sys
|
||||
import os
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Add project root to sys.path
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
# Mock comfy modules needed for import
|
||||
sys.modules['comfy'] = MagicMock()
|
||||
sys.modules['comfy.cli_args'] = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args = MagicMock()
|
||||
sys.modules['comfy.cli_args'].args.disable_metadata = False
|
||||
sys.modules['comfy.utils'] = MagicMock()
|
||||
sys.modules['folder_paths'] = MagicMock()
|
||||
sys.modules['server'] = MagicMock()
|
||||
|
||||
# Mock torch if not available (video_node imports it)
|
||||
if 'torch' not in sys.modules:
|
||||
sys.modules['torch'] = MagicMock()
|
||||
|
||||
from nodes.image_node import DiscordSendSaveImage
|
||||
from nodes.video_node import DiscordSendSaveVideo
|
||||
|
||||
class TestUXTooltips(unittest.TestCase):
|
||||
def test_video_node_add_time_tooltip(self):
|
||||
"""Test that the add_time tooltip in video node contains the critical warning."""
|
||||
input_types = DiscordSendSaveVideo.INPUT_TYPES()
|
||||
add_time_config = input_types["optional"]["add_time"]
|
||||
tooltip = add_time_config[1]["tooltip"]
|
||||
|
||||
# Verify it HAS the warning
|
||||
self.assertIn("CRITICAL", tooltip)
|
||||
self.assertIn("single-frame playback", tooltip)
|
||||
self.assertIn("Add time", tooltip)
|
||||
|
||||
def test_image_node_add_time_tooltip(self):
|
||||
"""Test that the add_time tooltip in image node is standard."""
|
||||
input_types = DiscordSendSaveImage.INPUT_TYPES()
|
||||
add_time_config = input_types["optional"]["add_time"]
|
||||
tooltip = add_time_config[1]["tooltip"]
|
||||
|
||||
expected = "Add time (HH-MM-SS) to the filename."
|
||||
self.assertEqual(tooltip, expected)
|
||||
|
||||
def test_github_token_tooltip(self):
|
||||
"""Test that the github_token tooltip contains helpful instructions."""
|
||||
# Both nodes inherit from BaseDiscordNode, so check one
|
||||
input_types = DiscordSendSaveImage.INPUT_TYPES()
|
||||
token_config = input_types["optional"]["github_token"]
|
||||
tooltip = token_config[1]["tooltip"]
|
||||
|
||||
self.assertIn("Settings > Developer settings > Tokens", tooltip)
|
||||
self.assertIn("Requires 'repo' scope", tooltip)
|
||||
|
||||
def test_resize_method_clarity(self):
|
||||
"""Test that resize_method tooltip clarifies dependency on resize_to_power_of_2."""
|
||||
input_types = DiscordSendSaveImage.INPUT_TYPES()
|
||||
resize_config = input_types["optional"]["resize_method"]
|
||||
tooltip = resize_config[1]["tooltip"]
|
||||
|
||||
self.assertIn("ONLY when 'resize_to_power_of_2' is enabled", tooltip)
|
||||
self.assertIn("Ignored otherwise", tooltip)
|
||||
self.assertIn("lanczos: Best for photos", tooltip)
|
||||
|
||||
def test_overwrite_safety_warning(self):
|
||||
"""Test that overwrite_last tooltip contains safety warning in both nodes."""
|
||||
# Test Image Node
|
||||
input_types_img = DiscordSendSaveImage.INPUT_TYPES()
|
||||
tooltip_img = input_types_img["required"]["overwrite_last"][1]["tooltip"]
|
||||
|
||||
self.assertIn("CAUTION", tooltip_img)
|
||||
self.assertIn("REPLACE the previous file", tooltip_img)
|
||||
self.assertIn("dangerous for batch production", tooltip_img)
|
||||
self.assertIn("disable 'add_time' and 'add_date'", tooltip_img)
|
||||
|
||||
# Test Video Node
|
||||
input_types_vid = DiscordSendSaveVideo.INPUT_TYPES()
|
||||
tooltip_vid = input_types_vid["required"]["overwrite_last"][1]["tooltip"]
|
||||
|
||||
self.assertIn("CAUTION", tooltip_vid)
|
||||
self.assertIn("REPLACE the previous file", tooltip_vid)
|
||||
self.assertIn("dangerous for batch production", tooltip_vid)
|
||||
self.assertIn("Disabling 'add_time' to overwrite files will cause single-frame playback issues", tooltip_vid)
|
||||
|
||||
def test_include_video_info_tooltip(self):
|
||||
"""Test that include_video_info tooltip provides the helpful tip about add_time."""
|
||||
input_types = DiscordSendSaveVideo.INPUT_TYPES()
|
||||
info_config = input_types["optional"]["include_video_info"]
|
||||
tooltip = info_config[1]["tooltip"]
|
||||
|
||||
self.assertIn("TIP", tooltip)
|
||||
self.assertIn("Disable this instead of 'add_time'", tooltip)
|
||||
self.assertIn("avoid the Discord single-frame bug", tooltip)
|
||||
|
||||
def test_resize_to_power_of_2_tooltip(self):
|
||||
"""Test that resize_to_power_of_2 tooltip warns about aspect ratio distortion."""
|
||||
input_types = DiscordSendSaveImage.INPUT_TYPES()
|
||||
resize_config = input_types["optional"]["resize_to_power_of_2"]
|
||||
tooltip = resize_config[1]["tooltip"]
|
||||
|
||||
self.assertIn("May distort aspect ratio", tooltip)
|
||||
self.assertIn("Uses the algorithm selected in 'resize_method'", tooltip)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,70 @@
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Mock modules that might be missing in the environment or causing issues
|
||||
sys.modules["folder_paths"] = MagicMock()
|
||||
sys.modules["comfy"] = MagicMock()
|
||||
sys.modules["server"] = MagicMock()
|
||||
sys.modules["torch"] = MagicMock()
|
||||
sys.modules["numpy"] = MagicMock()
|
||||
sys.modules["cv2"] = MagicMock()
|
||||
|
||||
# Import the module under test
|
||||
sys.path.append(os.getcwd())
|
||||
|
||||
from shared.discord.webhook_client import DiscordWebhookClient, send_to_discord_with_retry, sanitize_token_from_text
|
||||
|
||||
class TestWebhookSecurity(unittest.TestCase):
|
||||
def test_token_leak_in_client_error(self):
|
||||
"""
|
||||
Test that webhook tokens are NOT leaked in client error details.
|
||||
"""
|
||||
webhook_url = "https://discord.com/api/webhooks/123456789/SuperSecretToken123"
|
||||
client = DiscordWebhookClient(webhook_url)
|
||||
|
||||
# Mock response to simulate a 400 error that echoes the URL
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 400
|
||||
# Simulate an API that echoes the request URL in the error body
|
||||
mock_response.text = f"Error processing request to {webhook_url}: Invalid payload"
|
||||
mock_response.content = mock_response.text.encode('utf-8')
|
||||
|
||||
with patch('requests.post', return_value=mock_response):
|
||||
success, response = client.send_message("test")
|
||||
|
||||
self.assertFalse(success)
|
||||
error_details = response.get("details", "")
|
||||
|
||||
print(f"\nDEBUG: Error details: {error_details}")
|
||||
self.assertNotIn("SuperSecretToken123", error_details)
|
||||
self.assertIn("[REDACTED]", error_details)
|
||||
|
||||
def test_sanitize_token_from_text(self):
|
||||
"""
|
||||
Test the standalone sanitization function.
|
||||
"""
|
||||
webhook_url = "https://discord.com/api/webhooks/123456789/MySecretToken-Part2"
|
||||
|
||||
# Test 1: Simple URL in text
|
||||
text = f"Failed to send to {webhook_url}"
|
||||
sanitized = sanitize_token_from_text(text, webhook_url)
|
||||
self.assertNotIn("MySecretToken-Part2", sanitized)
|
||||
self.assertIn("[REDACTED]", sanitized)
|
||||
|
||||
# Test 2: Token embedded in other text
|
||||
text = "Some error occurred with token MySecretToken-Part2 processing"
|
||||
sanitized = sanitize_token_from_text(text, webhook_url)
|
||||
self.assertNotIn("MySecretToken-Part2", sanitized)
|
||||
self.assertIn("[REDACTED]", sanitized)
|
||||
|
||||
# Test 3: Multiple occurrences
|
||||
text = f"URL: {webhook_url}, Retry: {webhook_url}"
|
||||
sanitized = sanitize_token_from_text(text, webhook_url)
|
||||
self.assertNotIn("MySecretToken-Part2", sanitized)
|
||||
self.assertEqual(sanitized.count("[REDACTED]"), 2)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,183 @@
|
||||
import unittest
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from shared.workflow import WorkflowBuilder, extract_prompts_from_workflow
|
||||
|
||||
class TestWorkflowBuilder(unittest.TestCase):
|
||||
"""Tests for the WorkflowBuilder class."""
|
||||
|
||||
def setUp(self):
|
||||
# A simple standard ComfyUI API JSON format
|
||||
self.basic_api_json = {
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"seed": 123,
|
||||
"steps": 20,
|
||||
"cfg": 8,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1,
|
||||
"model": ["4", 0],
|
||||
"positive": ["6", 0],
|
||||
"negative": ["7", 0],
|
||||
"latent_image": ["5", 0]
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": "original positive prompt",
|
||||
"clip": ["4", 1]
|
||||
},
|
||||
"_meta": {"title": "Positive Prompt"}
|
||||
},
|
||||
"7": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {
|
||||
"text": "original negative prompt",
|
||||
"clip": ["4", 1]
|
||||
},
|
||||
"_meta": {"title": "Negative Prompt"}
|
||||
}
|
||||
}
|
||||
|
||||
def test_set_prompt_by_title(self):
|
||||
"""Should update prompts based on node titles."""
|
||||
builder = WorkflowBuilder(self.basic_api_json)
|
||||
builder.set_prompt("new positive", "new negative")
|
||||
|
||||
workflow = builder.get_workflow()
|
||||
self.assertEqual(workflow["6"]["inputs"]["text"], "new positive")
|
||||
self.assertEqual(workflow["7"]["inputs"]["text"], "new negative")
|
||||
|
||||
def test_set_seed(self):
|
||||
"""Should update seed in KSampler node."""
|
||||
builder = WorkflowBuilder(self.basic_api_json)
|
||||
builder.set_seed(999)
|
||||
|
||||
workflow = builder.get_workflow()
|
||||
self.assertEqual(workflow["3"]["inputs"]["seed"], 999)
|
||||
|
||||
def test_set_steps(self):
|
||||
"""Should update steps in KSampler node."""
|
||||
builder = WorkflowBuilder(self.basic_api_json)
|
||||
builder.set_steps(50)
|
||||
|
||||
workflow = builder.get_workflow()
|
||||
self.assertEqual(workflow["3"]["inputs"]["steps"], 50)
|
||||
|
||||
def test_set_cfg(self):
|
||||
"""Should update cfg in KSampler node."""
|
||||
builder = WorkflowBuilder(self.basic_api_json)
|
||||
builder.set_cfg(12.5)
|
||||
|
||||
workflow = builder.get_workflow()
|
||||
self.assertEqual(workflow["3"]["inputs"]["cfg"], 12.5)
|
||||
|
||||
def test_trace_input_fallback(self):
|
||||
"""Should update prompts by tracing KSampler inputs if titles are missing."""
|
||||
# Remove titles
|
||||
api_json = json.loads(json.dumps(self.basic_api_json))
|
||||
del api_json["6"]["_meta"]
|
||||
del api_json["7"]["_meta"]
|
||||
|
||||
builder = WorkflowBuilder(api_json)
|
||||
builder.set_prompt("traced positive", "traced negative")
|
||||
|
||||
workflow = builder.get_workflow()
|
||||
self.assertEqual(workflow["6"]["inputs"]["text"], "traced positive")
|
||||
self.assertEqual(workflow["7"]["inputs"]["text"], "traced negative")
|
||||
|
||||
class TestPromptExtractor(unittest.TestCase):
|
||||
"""Tests for the extract_prompts_from_workflow function."""
|
||||
|
||||
def test_extract_from_api_format(self):
|
||||
"""Should extract prompts from API-style JSON."""
|
||||
api_data = {
|
||||
"6": {
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["vibrant sunset over ocean"]
|
||||
},
|
||||
"7": {
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["low quality, blurry, watermark"]
|
||||
}
|
||||
}
|
||||
# Wrapped in a 'nodes' key for the extractor logic
|
||||
workflow = {"nodes": api_data}
|
||||
|
||||
pos, neg = extract_prompts_from_workflow(workflow)
|
||||
self.assertEqual(pos, "vibrant sunset over ocean")
|
||||
self.assertEqual(neg, "low quality, blurry, watermark")
|
||||
|
||||
def test_extract_heuristic_negative(self):
|
||||
"""Should correctly identify negative prompt based on content keywords."""
|
||||
api_data = {
|
||||
"1": {
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["ugly, deformed, mutated, extra fingers"]
|
||||
},
|
||||
"2": {
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["a majestic eagle"]
|
||||
}
|
||||
}
|
||||
workflow = {"nodes": api_data}
|
||||
|
||||
pos, neg = extract_prompts_from_workflow(workflow)
|
||||
self.assertEqual(pos, "a majestic eagle")
|
||||
self.assertEqual(neg, "ugly, deformed, mutated, extra fingers")
|
||||
|
||||
def test_extract_unrelated_text_nodes(self):
|
||||
"""Should ignore empty strings and pick lowest score as positive prompt."""
|
||||
workflow = {
|
||||
"6": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "embedding:bad-artist, EasyNegative, low quality"} # Score: 3
|
||||
},
|
||||
"7": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": " "} # Empty/Whitespace - Should be ignored
|
||||
},
|
||||
"8": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "v1-5-pruned-emaonly.safetensors"} # Unrelated - Score: 0 (or low)
|
||||
},
|
||||
"9": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "a beautiful forest, cinematic lighting"} # Target - Score: 0
|
||||
}
|
||||
}
|
||||
|
||||
# In this case, node 8 and 9 both have score 0.
|
||||
# The logic picks the last one in the list (node 9).
|
||||
# Node 7 should be completely ignored.
|
||||
|
||||
# We need to ensure extraction finds the right pair
|
||||
positive, negative = extract_prompts_from_workflow(workflow)
|
||||
|
||||
self.assertEqual(negative, "embedding:bad-artist, EasyNegative, low quality")
|
||||
self.assertEqual(positive, "a beautiful forest, cinematic lighting")
|
||||
|
||||
def test_extract_single_prompt(self):
|
||||
"""Should assume single CLIP node is positive."""
|
||||
api_data = {
|
||||
"1": {
|
||||
"type": "CLIPTextEncode",
|
||||
"widgets_values": ["solo traveler"]
|
||||
}
|
||||
}
|
||||
workflow = {"nodes": api_data}
|
||||
|
||||
pos, neg = extract_prompts_from_workflow(workflow)
|
||||
self.assertEqual(pos, "solo traveler")
|
||||
self.assertEqual(neg, "")
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,14 +0,0 @@
|
||||
{
|
||||
"clones": {
|
||||
"count": 211,
|
||||
"uniques": 166
|
||||
},
|
||||
"views": {
|
||||
"count": 19,
|
||||
"uniques": 18
|
||||
},
|
||||
"downloads": {
|
||||
"smart_count": 0
|
||||
},
|
||||
"timestamp": "2026-04-21T02:13:13.650Z"
|
||||
}
|
||||
@@ -0,0 +1,732 @@
|
||||
{
|
||||
"last_node_id": 69,
|
||||
"last_link_id": 117,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 7,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-810,
|
||||
-2740
|
||||
],
|
||||
"size": [
|
||||
425.27801513671875,
|
||||
180.6060791015625
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 75
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
101
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"title": "CLIP Text Encode (Negative Prompt)",
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Overexposure, static, blurred details, subtitles, paintings, pictures, still, overall gray, worst quality, low quality, JPEG compression residue, ugly, mutilated, redundant fingers, poorly painted hands, poorly painted faces, deformed, disfigured, deformed limbs, fused fingers, cluttered background, three legs, a lot of people in the background, upside down"
|
||||
],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VAEDecode",
|
||||
"pos": [
|
||||
100,
|
||||
-2430
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "samples",
|
||||
"type": "LATENT",
|
||||
"link": 35
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": 76
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
117
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAEDecode"
|
||||
},
|
||||
"widgets_values": [],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 38,
|
||||
"type": "CLIPLoader",
|
||||
"pos": [
|
||||
-1250,
|
||||
-2660
|
||||
],
|
||||
"size": [
|
||||
390,
|
||||
82
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"links": [
|
||||
74,
|
||||
75
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"umt5_xxl_fp8_e4m3fn_scaled.safetensors",
|
||||
"wan",
|
||||
"default"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 53,
|
||||
"type": "CLIPVisionLoader",
|
||||
"pos": [
|
||||
-1180,
|
||||
-2530
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP_VISION",
|
||||
"type": "CLIP_VISION",
|
||||
"links": [
|
||||
106
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPVisionLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"clip_vision_h.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 54,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
-1180,
|
||||
-2420
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
108,
|
||||
109
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Img_01095_.png",
|
||||
"image"
|
||||
],
|
||||
"color": "#223",
|
||||
"bgcolor": "#335"
|
||||
},
|
||||
{
|
||||
"id": 51,
|
||||
"type": "WanImageToVideo",
|
||||
"pos": [
|
||||
-310,
|
||||
-2960
|
||||
],
|
||||
"size": [
|
||||
342.5999755859375,
|
||||
210
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"link": 100
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"link": 101
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": 105
|
||||
},
|
||||
{
|
||||
"name": "clip_vision_output",
|
||||
"type": "CLIP_VISION_OUTPUT",
|
||||
"shape": 7,
|
||||
"link": 107
|
||||
},
|
||||
{
|
||||
"name": "start_image",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": 109
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
102
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
103
|
||||
],
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "latent",
|
||||
"type": "LATENT",
|
||||
"links": [
|
||||
104
|
||||
],
|
||||
"slot_index": 2
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "WanImageToVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
512,
|
||||
512,
|
||||
33,
|
||||
1
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 52,
|
||||
"type": "CLIPVisionEncode",
|
||||
"pos": [
|
||||
-810,
|
||||
-2490
|
||||
],
|
||||
"size": [
|
||||
380.4000244140625,
|
||||
78
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip_vision",
|
||||
"type": "CLIP_VISION",
|
||||
"link": 106
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 108
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP_VISION_OUTPUT",
|
||||
"type": "CLIP_VISION_OUTPUT",
|
||||
"links": [
|
||||
107
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPVisionEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"none"
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-810,
|
||||
-2960
|
||||
],
|
||||
"size": [
|
||||
422.84503173828125,
|
||||
164.31304931640625
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 74
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
100
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"title": "CLIP Text Encode (Positive Prompt)",
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
""
|
||||
],
|
||||
"color": "#232",
|
||||
"bgcolor": "#353"
|
||||
},
|
||||
{
|
||||
"id": 39,
|
||||
"type": "VAELoader",
|
||||
"pos": [
|
||||
-1170,
|
||||
-2780
|
||||
],
|
||||
"size": [
|
||||
306.36004638671875,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"links": [
|
||||
76,
|
||||
105
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAELoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"wan_2.1_vae.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 55,
|
||||
"type": "UnetLoaderGGUFAdvanced",
|
||||
"pos": [
|
||||
-1200,
|
||||
-2960
|
||||
],
|
||||
"size": [
|
||||
340.20001220703125,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"links": [
|
||||
110
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "city96/ComfyUI-GGUF",
|
||||
"ver": "298192ed60f8ca821c6fe5f8030cae23424cada5",
|
||||
"Node name for S&R": "UnetLoaderGGUFAdvanced"
|
||||
},
|
||||
"widgets_values": [
|
||||
"wan2.1-i2v-14b-480p-Q4_0.gguf",
|
||||
"default",
|
||||
"default",
|
||||
false
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "KSampler",
|
||||
"pos": [
|
||||
90,
|
||||
-2960
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
474
|
||||
],
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 110
|
||||
},
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"link": 102
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"link": 103
|
||||
},
|
||||
{
|
||||
"name": "latent_image",
|
||||
"type": "LATENT",
|
||||
"link": 104
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "LATENT",
|
||||
"type": "LATENT",
|
||||
"links": [
|
||||
35
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "KSampler"
|
||||
},
|
||||
"widgets_values": [
|
||||
676833766536034,
|
||||
"fixed",
|
||||
10,
|
||||
6,
|
||||
"dpmpp_2m",
|
||||
"sgm_uniform",
|
||||
1
|
||||
],
|
||||
"color": "#233",
|
||||
"bgcolor": "#355"
|
||||
},
|
||||
{
|
||||
"id": 69,
|
||||
"type": "DiscordSendSaveVideo",
|
||||
"pos": [
|
||||
490,
|
||||
-2960
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
504
|
||||
],
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 117
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "AUDIO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "video_path",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "AEmotionStudio/ComfyUI-DiscordSend",
|
||||
"ver": "3b2e21e1d3b90d04bd54bd51e9b5b29d405f2c17",
|
||||
"Node name for S&R": "DiscordSendSaveVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"ComfyUI-Video",
|
||||
"disable",
|
||||
"video/h264-mp4",
|
||||
8,
|
||||
85,
|
||||
0,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
"enable",
|
||||
"enable",
|
||||
"enable",
|
||||
true,
|
||||
"",
|
||||
"",
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
"color": "#222",
|
||||
"bgcolor": "#000"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
35,
|
||||
3,
|
||||
0,
|
||||
8,
|
||||
0,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
74,
|
||||
38,
|
||||
0,
|
||||
6,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
75,
|
||||
38,
|
||||
0,
|
||||
7,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
76,
|
||||
39,
|
||||
0,
|
||||
8,
|
||||
1,
|
||||
"VAE"
|
||||
],
|
||||
[
|
||||
100,
|
||||
6,
|
||||
0,
|
||||
51,
|
||||
0,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
101,
|
||||
7,
|
||||
0,
|
||||
51,
|
||||
1,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
102,
|
||||
51,
|
||||
0,
|
||||
3,
|
||||
1,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
103,
|
||||
51,
|
||||
1,
|
||||
3,
|
||||
2,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
104,
|
||||
51,
|
||||
2,
|
||||
3,
|
||||
3,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
105,
|
||||
39,
|
||||
0,
|
||||
51,
|
||||
2,
|
||||
"VAE"
|
||||
],
|
||||
[
|
||||
106,
|
||||
53,
|
||||
0,
|
||||
52,
|
||||
0,
|
||||
"CLIP_VISION"
|
||||
],
|
||||
[
|
||||
107,
|
||||
52,
|
||||
0,
|
||||
51,
|
||||
3,
|
||||
"CLIP_VISION_OUTPUT"
|
||||
],
|
||||
[
|
||||
108,
|
||||
54,
|
||||
0,
|
||||
52,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
109,
|
||||
54,
|
||||
0,
|
||||
51,
|
||||
4,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
110,
|
||||
55,
|
||||
0,
|
||||
3,
|
||||
0,
|
||||
"MODEL"
|
||||
],
|
||||
[
|
||||
117,
|
||||
8,
|
||||
0,
|
||||
69,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.6852083333333341,
|
||||
"offset": [
|
||||
1492.0798316376156,
|
||||
3217.6726640627494
|
||||
]
|
||||
},
|
||||
"VHS_latentpreview": false,
|
||||
"VHS_latentpreviewrate": 0,
|
||||
"VHS_MetadataImage": true,
|
||||
"VHS_KeepIntermediate": true
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,527 @@
|
||||
{
|
||||
"last_node_id": 90,
|
||||
"last_link_id": 143,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 70,
|
||||
"type": "CheckpointLoaderSimple",
|
||||
"pos": [
|
||||
-1193.2313232421875,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
98
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
143
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
118
|
||||
],
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"shape": 3,
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CheckpointLoaderSimple"
|
||||
},
|
||||
"widgets_values": [
|
||||
"dreamshaper_8.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 71,
|
||||
"type": "VAELoader",
|
||||
"pos": [
|
||||
-1264.126220703125,
|
||||
-2002.368896484375
|
||||
],
|
||||
"size": [
|
||||
385.8948669433594,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
120
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAELoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"SD1.5\\vae-ft-mse-840000-ema-pruned.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 74,
|
||||
"type": "VAEDecode",
|
||||
"pos": [
|
||||
-300,
|
||||
-1600
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "samples",
|
||||
"type": "LATENT",
|
||||
"link": 119
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": 120
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
142
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAEDecode"
|
||||
},
|
||||
"widgets_values": [],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 85,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-799.4965209960938,
|
||||
-1967.837890625
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 135
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
137
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"bad, ugly, two tails, malformed, deformed, horror, (worst quality, low quality: 1.4),"
|
||||
],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 84,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-799.4965209960938,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 134
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
136
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"a majestic white tiger sitting near a waterfall, center view, symmetrical, flowing water, white lillies, white flowers, nature, moss, beautiful trees, lush forest, serene, calming, high quality, masterpiece, "
|
||||
],
|
||||
"color": "#232",
|
||||
"bgcolor": "#353"
|
||||
},
|
||||
{
|
||||
"id": 72,
|
||||
"type": "CLIPSetLastLayer",
|
||||
"pos": [
|
||||
-1193.2313232421875,
|
||||
-2117
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 118
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
134,
|
||||
135
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPSetLastLayer"
|
||||
},
|
||||
"widgets_values": [
|
||||
-2
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 73,
|
||||
"type": "EmptyLatentImage",
|
||||
"pos": [
|
||||
-300,
|
||||
-1750
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "LATENT",
|
||||
"type": "LATENT",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
125
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "EmptyLatentImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
768,
|
||||
768,
|
||||
1
|
||||
],
|
||||
"color": "#2a363b",
|
||||
"bgcolor": "#3f5159"
|
||||
},
|
||||
{
|
||||
"id": 76,
|
||||
"type": "KSampler",
|
||||
"pos": [
|
||||
-300,
|
||||
-2270
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
474
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 143
|
||||
},
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"link": 136
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"link": 137
|
||||
},
|
||||
{
|
||||
"name": "latent_image",
|
||||
"type": "LATENT",
|
||||
"link": 125
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "LATENT",
|
||||
"type": "LATENT",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
119
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "KSampler"
|
||||
},
|
||||
"widgets_values": [
|
||||
888937,
|
||||
"fixed",
|
||||
30,
|
||||
8,
|
||||
"uni_pc",
|
||||
"karras",
|
||||
1
|
||||
],
|
||||
"color": "#233",
|
||||
"bgcolor": "#355"
|
||||
},
|
||||
{
|
||||
"id": 90,
|
||||
"type": "DiscordSendSaveImage",
|
||||
"pos": [
|
||||
100,
|
||||
-2270
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
508
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 142
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image_path",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "AEmotionStudio/ComfyUI-DiscordSend",
|
||||
"ver": "3b2e21e1d3b90d04bd54bd51e9b5b29d405f2c17",
|
||||
"Node name for S&R": "DiscordSendSaveImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"ComfyUI-Image",
|
||||
"disable",
|
||||
"png",
|
||||
95,
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
"enable",
|
||||
"enable",
|
||||
"enable",
|
||||
"disable",
|
||||
"lanczos",
|
||||
true,
|
||||
"",
|
||||
"",
|
||||
true,
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
"color": "#222",
|
||||
"bgcolor": "#000"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
118,
|
||||
70,
|
||||
1,
|
||||
72,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
119,
|
||||
76,
|
||||
0,
|
||||
74,
|
||||
0,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
120,
|
||||
71,
|
||||
0,
|
||||
74,
|
||||
1,
|
||||
"VAE"
|
||||
],
|
||||
[
|
||||
125,
|
||||
73,
|
||||
0,
|
||||
76,
|
||||
3,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
134,
|
||||
72,
|
||||
0,
|
||||
84,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
135,
|
||||
72,
|
||||
0,
|
||||
85,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
136,
|
||||
84,
|
||||
0,
|
||||
76,
|
||||
1,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
137,
|
||||
85,
|
||||
0,
|
||||
76,
|
||||
2,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
142,
|
||||
74,
|
||||
0,
|
||||
90,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
143,
|
||||
70,
|
||||
0,
|
||||
76,
|
||||
0,
|
||||
"MODEL"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.8264462809917354,
|
||||
"offset": [
|
||||
1306.035455615458,
|
||||
2467.929546393067
|
||||
]
|
||||
},
|
||||
"VHS_latentpreview": false,
|
||||
"VHS_latentpreviewrate": 0,
|
||||
"VHS_MetadataImage": true,
|
||||
"VHS_KeepIntermediate": true
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,849 @@
|
||||
{
|
||||
"last_node_id": 91,
|
||||
"last_link_id": 145,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 71,
|
||||
"type": "VAELoader",
|
||||
"pos": [
|
||||
-1264.126220703125,
|
||||
-2002.368896484375
|
||||
],
|
||||
"size": [
|
||||
385.8948669433594,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
120
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAELoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"SD1.5\\vae-ft-mse-840000-ema-pruned.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 75,
|
||||
"type": "ADE_AnimateDiffLoaderGen1",
|
||||
"pos": [
|
||||
-301.136474609375,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
242
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 145
|
||||
},
|
||||
{
|
||||
"name": "context_options",
|
||||
"type": "CONTEXT_OPTIONS",
|
||||
"shape": 7,
|
||||
"link": 140
|
||||
},
|
||||
{
|
||||
"name": "motion_lora",
|
||||
"type": "MOTION_LORA",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "ad_settings",
|
||||
"type": "AD_SETTINGS",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "ad_keyframes",
|
||||
"type": "AD_KEYFRAMES",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "sample_settings",
|
||||
"type": "SAMPLE_SETTINGS",
|
||||
"shape": 7,
|
||||
"link": 141
|
||||
},
|
||||
{
|
||||
"name": "scale_multival",
|
||||
"type": "MULTIVAL",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "effect_multival",
|
||||
"type": "MULTIVAL",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "per_block",
|
||||
"type": "PER_BLOCK",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"links": [
|
||||
144
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfyui-animatediff-evolved",
|
||||
"ver": "916dd55092f4375f0e566e27515959e465c8e0df",
|
||||
"Node name for S&R": "ADE_AnimateDiffLoaderGen1"
|
||||
},
|
||||
"widgets_values": [
|
||||
"v3_sd15_mm.ckpt",
|
||||
"sqrt_linear (AnimateDiff)"
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 87,
|
||||
"type": "ADE_StandardStaticContextOptions",
|
||||
"pos": [
|
||||
-301.136474609375,
|
||||
-1971.269775390625
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "prev_context",
|
||||
"type": "CONTEXT_OPTIONS",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "view_opts",
|
||||
"type": "VIEW_OPTS",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONTEXT_OPTS",
|
||||
"type": "CONTEXT_OPTIONS",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
140
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfyui-animatediff-evolved",
|
||||
"ver": "916dd55092f4375f0e566e27515959e465c8e0df",
|
||||
"Node name for S&R": "ADE_StandardStaticContextOptions"
|
||||
},
|
||||
"widgets_values": [
|
||||
16,
|
||||
4,
|
||||
"pyramid",
|
||||
false,
|
||||
0,
|
||||
1
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 88,
|
||||
"type": "ADE_AnimateDiffSamplingSettings",
|
||||
"pos": [
|
||||
-301.136474609375,
|
||||
-1718.2647705078125
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
294
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "noise_layers",
|
||||
"type": "NOISE_LAYERS",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "iteration_opts",
|
||||
"type": "ITERATION_OPTS",
|
||||
"shape": 7,
|
||||
"link": 139,
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "custom_cfg",
|
||||
"type": "CUSTOM_CFG",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "sigma_schedule",
|
||||
"type": "SIGMA_SCHEDULE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "seed_override",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed_override"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "ancestral_opts",
|
||||
"type": "ANCESTRAL_OPTS",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_inject",
|
||||
"type": "IMAGE_INJECT",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_inject",
|
||||
"type": "IMAGE_INJECT",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "settings",
|
||||
"type": "SAMPLE_SETTINGS",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
141
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfyui-animatediff-evolved",
|
||||
"ver": "916dd55092f4375f0e566e27515959e465c8e0df",
|
||||
"Node name for S&R": "ADE_AnimateDiffSamplingSettings"
|
||||
},
|
||||
"widgets_values": [
|
||||
0,
|
||||
"FreeNoise",
|
||||
"comfy",
|
||||
0,
|
||||
0,
|
||||
false
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 85,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-799.4965209960938,
|
||||
-1967.837890625
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 135
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
137
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"bad, ugly, two tails, malformed, deformed, horror, (worst quality, low quality: 1.4),"
|
||||
],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 84,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
-799.4965209960938,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 134
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
136
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"a majestic white tiger sitting near a waterfall, center view, symmetrical, flowing water, white lillies, white flowers, nature, moss, beautiful trees, lush forest, serene, calming, high quality, masterpiece, "
|
||||
],
|
||||
"color": "#232",
|
||||
"bgcolor": "#353"
|
||||
},
|
||||
{
|
||||
"id": 89,
|
||||
"type": "ADE_IterationOptsFreeInit",
|
||||
"pos": [
|
||||
-301.136474609375,
|
||||
-1383.185791015625
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
274
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "ITERATION_OPTS",
|
||||
"type": "ITERATION_OPTS",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
139
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfyui-animatediff-evolved",
|
||||
"ver": "916dd55092f4375f0e566e27515959e465c8e0df",
|
||||
"Node name for S&R": "ADE_IterationOptsFreeInit"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
"gaussian",
|
||||
0.25,
|
||||
0.25,
|
||||
4,
|
||||
999,
|
||||
false,
|
||||
"FreeInit [sampler sigma]",
|
||||
0,
|
||||
1
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 72,
|
||||
"type": "CLIPSetLastLayer",
|
||||
"pos": [
|
||||
-1193.2313232421875,
|
||||
-2117
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 118
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
134,
|
||||
135
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CLIPSetLastLayer"
|
||||
},
|
||||
"widgets_values": [
|
||||
-2
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#653"
|
||||
},
|
||||
{
|
||||
"id": 74,
|
||||
"type": "VAEDecode",
|
||||
"pos": [
|
||||
106.84590911865234,
|
||||
-1595.5885009765625
|
||||
],
|
||||
"size": [
|
||||
210,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "samples",
|
||||
"type": "LATENT",
|
||||
"link": 119
|
||||
},
|
||||
{
|
||||
"name": "vae",
|
||||
"type": "VAE",
|
||||
"link": 120
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
143
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "VAEDecode"
|
||||
},
|
||||
"widgets_values": [],
|
||||
"color": "#322",
|
||||
"bgcolor": "#533"
|
||||
},
|
||||
{
|
||||
"id": 70,
|
||||
"type": "CheckpointLoaderSimple",
|
||||
"pos": [
|
||||
-1193.2313232421875,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
98
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MODEL",
|
||||
"type": "MODEL",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
145
|
||||
],
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
118
|
||||
],
|
||||
"slot_index": 1
|
||||
},
|
||||
{
|
||||
"name": "VAE",
|
||||
"type": "VAE",
|
||||
"shape": 3,
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "CheckpointLoaderSimple"
|
||||
},
|
||||
"widgets_values": [
|
||||
"dreamshaper_8.safetensors"
|
||||
],
|
||||
"color": "#332922",
|
||||
"bgcolor": "#593930"
|
||||
},
|
||||
{
|
||||
"id": 76,
|
||||
"type": "KSampler",
|
||||
"pos": [
|
||||
106.22499084472656,
|
||||
-2268.5
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
474
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL",
|
||||
"link": 144
|
||||
},
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"link": 136
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"link": 137
|
||||
},
|
||||
{
|
||||
"name": "latent_image",
|
||||
"type": "LATENT",
|
||||
"link": 125
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "LATENT",
|
||||
"type": "LATENT",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
119
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "KSampler"
|
||||
},
|
||||
"widgets_values": [
|
||||
888936,
|
||||
"fixed",
|
||||
30,
|
||||
8,
|
||||
"uni_pc",
|
||||
"karras",
|
||||
1
|
||||
],
|
||||
"color": "#233",
|
||||
"bgcolor": "#355"
|
||||
},
|
||||
{
|
||||
"id": 73,
|
||||
"type": "EmptyLatentImage",
|
||||
"pos": [
|
||||
106.22499084472656,
|
||||
-1750
|
||||
],
|
||||
"size": [
|
||||
315,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "LATENT",
|
||||
"type": "LATENT",
|
||||
"shape": 3,
|
||||
"links": [
|
||||
125
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"cnr_id": "comfy-core",
|
||||
"ver": "0.3.18",
|
||||
"Node name for S&R": "EmptyLatentImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
768,
|
||||
768,
|
||||
16
|
||||
],
|
||||
"color": "#2a363b",
|
||||
"bgcolor": "#3f5159"
|
||||
},
|
||||
{
|
||||
"id": 91,
|
||||
"type": "DiscordSendSaveVideo",
|
||||
"pos": [
|
||||
503.32611083984375,
|
||||
-2268.85009765625
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
504
|
||||
],
|
||||
"flags": {},
|
||||
"order": 12,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 143
|
||||
},
|
||||
{
|
||||
"name": "audio",
|
||||
"type": "AUDIO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "video_path",
|
||||
"type": "STRING",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"aux_id": "AEmotionStudio/ComfyUI-DiscordSend",
|
||||
"ver": "3b2e21e1d3b90d04bd54bd51e9b5b29d405f2c17",
|
||||
"Node name for S&R": "DiscordSendSaveVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"ComfyUI-Video",
|
||||
"disable",
|
||||
"video/h264-mp4",
|
||||
8,
|
||||
85,
|
||||
0,
|
||||
false,
|
||||
false,
|
||||
true,
|
||||
"enable",
|
||||
"enable",
|
||||
"enable",
|
||||
true,
|
||||
"",
|
||||
"",
|
||||
true,
|
||||
true,
|
||||
true
|
||||
],
|
||||
"color": "#222",
|
||||
"bgcolor": "#000"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
118,
|
||||
70,
|
||||
1,
|
||||
72,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
119,
|
||||
76,
|
||||
0,
|
||||
74,
|
||||
0,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
120,
|
||||
71,
|
||||
0,
|
||||
74,
|
||||
1,
|
||||
"VAE"
|
||||
],
|
||||
[
|
||||
125,
|
||||
73,
|
||||
0,
|
||||
76,
|
||||
3,
|
||||
"LATENT"
|
||||
],
|
||||
[
|
||||
134,
|
||||
72,
|
||||
0,
|
||||
84,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
135,
|
||||
72,
|
||||
0,
|
||||
85,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
136,
|
||||
84,
|
||||
0,
|
||||
76,
|
||||
1,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
137,
|
||||
85,
|
||||
0,
|
||||
76,
|
||||
2,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
139,
|
||||
89,
|
||||
0,
|
||||
88,
|
||||
1,
|
||||
"ITERATION_OPTS"
|
||||
],
|
||||
[
|
||||
140,
|
||||
87,
|
||||
0,
|
||||
75,
|
||||
1,
|
||||
"CONTEXT_OPTIONS"
|
||||
],
|
||||
[
|
||||
141,
|
||||
88,
|
||||
0,
|
||||
75,
|
||||
5,
|
||||
"SAMPLE_SETTINGS"
|
||||
],
|
||||
[
|
||||
143,
|
||||
74,
|
||||
0,
|
||||
91,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
144,
|
||||
75,
|
||||
0,
|
||||
76,
|
||||
0,
|
||||
"MODEL"
|
||||
],
|
||||
[
|
||||
145,
|
||||
70,
|
||||
0,
|
||||
75,
|
||||
0,
|
||||
"MODEL"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.6209213230591553,
|
||||
"offset": [
|
||||
1429.8822830700044,
|
||||
2462.6955046657895
|
||||
]
|
||||
},
|
||||
"VHS_latentpreview": false,
|
||||
"VHS_latentpreviewrate": 0,
|
||||
"VHS_MetadataImage": true,
|
||||
"VHS_KeepIntermediate": true
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
Reference in New Issue
Block a user