vibe-aigc 0.3.0__tar.gz → 0.4.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {vibe_aigc-0.3.0/vibe_aigc.egg-info → vibe_aigc-0.4.0}/PKG-INFO +1 -1
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/pyproject.toml +2 -1
- vibe_aigc-0.4.0/vibe_aigc/character.py +457 -0
- vibe_aigc-0.4.0/vibe_aigc/video.py +388 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0/vibe_aigc.egg-info}/PKG-INFO +1 -1
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc.egg-info/SOURCES.txt +2 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/LICENSE +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/README.md +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/setup.cfg +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_adaptive_replanning.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_agents.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_assets.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_auto_checkpoint.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_automatic_checkpoints.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_checkpoint_serialization.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_error_handling.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_executor.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_feedback_system.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_integration.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_knowledge_base.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_metaplanner_resume.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_metaplanner_visualization.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_models.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_parallel_execution.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_planner.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_progress_callbacks.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_tools.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_visualization.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/tests/test_workflow_resume.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/__init__.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/agents.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/assets.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/cli.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/comfyui.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/executor.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/knowledge.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/llm.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/models.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/persistence.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/planner.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/tools.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/tools_multimodal.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc/visualization.py +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc.egg-info/dependency_links.txt +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc.egg-info/entry_points.txt +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc.egg-info/requires.txt +0 -0
- {vibe_aigc-0.3.0 → vibe_aigc-0.4.0}/vibe_aigc.egg-info/top_level.txt +0 -0
|
@@ -8,7 +8,7 @@ exclude = ["tests*", "docs*", "examples*", "landing*"]
|
|
|
8
8
|
|
|
9
9
|
[project]
|
|
10
10
|
name = "vibe-aigc"
|
|
11
|
-
version = "0.
|
|
11
|
+
version = "0.4.0"
|
|
12
12
|
description = "A New Paradigm for Content Generation via Agentic Orchestration"
|
|
13
13
|
authors = [{name = "Vibe AIGC Contributors"}]
|
|
14
14
|
license = "MIT"
|
|
@@ -68,3 +68,4 @@ warn_unused_configs = true
|
|
|
68
68
|
ignore_missing_imports = true
|
|
69
69
|
|
|
70
70
|
|
|
71
|
+
|
|
@@ -0,0 +1,457 @@
|
|
|
1
|
+
"""Character consistency using IPAdapter.
|
|
2
|
+
|
|
3
|
+
This enables maintaining character appearance across multiple
|
|
4
|
+
generations - critical for music videos, stories, etc.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import asyncio
|
|
8
|
+
import json
|
|
9
|
+
import uuid
|
|
10
|
+
import aiohttp
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
from typing import Any, Dict, List, Optional
|
|
13
|
+
from dataclasses import dataclass, field
|
|
14
|
+
|
|
15
|
+
from .comfyui import ComfyUIConfig, GenerationResult
|
|
16
|
+
from .tools import BaseTool, ToolResult, ToolSpec, ToolCategory
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class CharacterReference:
|
|
21
|
+
"""A character reference for consistency."""
|
|
22
|
+
name: str
|
|
23
|
+
reference_image: str # Path or URL to reference image
|
|
24
|
+
description: str = ""
|
|
25
|
+
weight: float = 0.8 # How strongly to apply the reference (0-1)
|
|
26
|
+
|
|
27
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
28
|
+
return {
|
|
29
|
+
"name": self.name,
|
|
30
|
+
"reference_image": self.reference_image,
|
|
31
|
+
"description": self.description,
|
|
32
|
+
"weight": self.weight
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class CharacterBank:
|
|
37
|
+
"""Bank of character references for consistency.
|
|
38
|
+
|
|
39
|
+
Like the paper's Character Bank - maintains character identity
|
|
40
|
+
across multiple generations.
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
def __init__(self):
|
|
44
|
+
self._characters: Dict[str, CharacterReference] = {}
|
|
45
|
+
|
|
46
|
+
def add(self, character: CharacterReference) -> None:
|
|
47
|
+
"""Add a character to the bank."""
|
|
48
|
+
self._characters[character.name.lower()] = character
|
|
49
|
+
|
|
50
|
+
def get(self, name: str) -> Optional[CharacterReference]:
|
|
51
|
+
"""Get a character by name."""
|
|
52
|
+
return self._characters.get(name.lower())
|
|
53
|
+
|
|
54
|
+
def list_characters(self) -> List[str]:
|
|
55
|
+
"""List all character names."""
|
|
56
|
+
return list(self._characters.keys())
|
|
57
|
+
|
|
58
|
+
def remove(self, name: str) -> bool:
|
|
59
|
+
"""Remove a character from the bank."""
|
|
60
|
+
if name.lower() in self._characters:
|
|
61
|
+
del self._characters[name.lower()]
|
|
62
|
+
return True
|
|
63
|
+
return False
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class IPAdapterBackend:
|
|
67
|
+
"""Backend for character-consistent generation using IPAdapter."""
|
|
68
|
+
|
|
69
|
+
def __init__(self, config: Optional[ComfyUIConfig] = None):
|
|
70
|
+
self.config = config or ComfyUIConfig()
|
|
71
|
+
self._client_id = str(uuid.uuid4())
|
|
72
|
+
|
|
73
|
+
async def is_available(self) -> bool:
|
|
74
|
+
"""Check if IPAdapter nodes are available."""
|
|
75
|
+
try:
|
|
76
|
+
async with aiohttp.ClientSession() as session:
|
|
77
|
+
async with session.get(
|
|
78
|
+
f"{self.config.base_url}/object_info",
|
|
79
|
+
timeout=aiohttp.ClientTimeout(total=5)
|
|
80
|
+
) as resp:
|
|
81
|
+
if resp.status != 200:
|
|
82
|
+
return False
|
|
83
|
+
obj_info = await resp.json()
|
|
84
|
+
# Check for IPAdapter nodes
|
|
85
|
+
return "IPAdapterApply" in obj_info or "IPAdapter" in obj_info
|
|
86
|
+
except Exception:
|
|
87
|
+
return False
|
|
88
|
+
|
|
89
|
+
async def generate_with_reference(
|
|
90
|
+
self,
|
|
91
|
+
prompt: str,
|
|
92
|
+
reference_image: str,
|
|
93
|
+
negative_prompt: str = "",
|
|
94
|
+
width: int = 512,
|
|
95
|
+
height: int = 512,
|
|
96
|
+
steps: int = 20,
|
|
97
|
+
cfg: float = 7.0,
|
|
98
|
+
seed: Optional[int] = None,
|
|
99
|
+
reference_weight: float = 0.8,
|
|
100
|
+
checkpoint: str = "v1-5-pruned-emaonly.safetensors"
|
|
101
|
+
) -> GenerationResult:
|
|
102
|
+
"""Generate an image consistent with a reference.
|
|
103
|
+
|
|
104
|
+
Args:
|
|
105
|
+
prompt: Text description
|
|
106
|
+
reference_image: Path to reference image for consistency
|
|
107
|
+
negative_prompt: What to avoid
|
|
108
|
+
width: Output width
|
|
109
|
+
height: Output height
|
|
110
|
+
steps: Sampling steps
|
|
111
|
+
cfg: Guidance scale
|
|
112
|
+
seed: Random seed
|
|
113
|
+
reference_weight: How strongly to apply reference (0-1)
|
|
114
|
+
checkpoint: SD checkpoint
|
|
115
|
+
|
|
116
|
+
Returns:
|
|
117
|
+
GenerationResult with consistent image
|
|
118
|
+
"""
|
|
119
|
+
if seed is None:
|
|
120
|
+
import random
|
|
121
|
+
seed = random.randint(0, 2**32 - 1)
|
|
122
|
+
|
|
123
|
+
# First, upload the reference image if it's a local path
|
|
124
|
+
if not reference_image.startswith(('http://', 'https://')):
|
|
125
|
+
reference_image = await self._upload_image(reference_image)
|
|
126
|
+
|
|
127
|
+
workflow = self._build_ipadapter_workflow(
|
|
128
|
+
prompt=prompt,
|
|
129
|
+
reference_image=reference_image,
|
|
130
|
+
negative_prompt=negative_prompt,
|
|
131
|
+
width=width,
|
|
132
|
+
height=height,
|
|
133
|
+
steps=steps,
|
|
134
|
+
cfg=cfg,
|
|
135
|
+
seed=seed,
|
|
136
|
+
reference_weight=reference_weight,
|
|
137
|
+
checkpoint=checkpoint
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
return await self._execute_workflow(workflow)
|
|
141
|
+
|
|
142
|
+
async def _upload_image(self, image_path: str) -> str:
|
|
143
|
+
"""Upload a local image to ComfyUI."""
|
|
144
|
+
path = Path(image_path)
|
|
145
|
+
if not path.exists():
|
|
146
|
+
raise FileNotFoundError(f"Reference image not found: {image_path}")
|
|
147
|
+
|
|
148
|
+
async with aiohttp.ClientSession() as session:
|
|
149
|
+
data = aiohttp.FormData()
|
|
150
|
+
data.add_field('image',
|
|
151
|
+
open(path, 'rb'),
|
|
152
|
+
filename=path.name,
|
|
153
|
+
content_type='image/png')
|
|
154
|
+
|
|
155
|
+
async with session.post(
|
|
156
|
+
f"{self.config.base_url}/upload/image",
|
|
157
|
+
data=data
|
|
158
|
+
) as resp:
|
|
159
|
+
if resp.status != 200:
|
|
160
|
+
raise RuntimeError(f"Failed to upload image: {await resp.text()}")
|
|
161
|
+
result = await resp.json()
|
|
162
|
+
return result.get("name", path.name)
|
|
163
|
+
|
|
164
|
+
def _build_ipadapter_workflow(
|
|
165
|
+
self,
|
|
166
|
+
prompt: str,
|
|
167
|
+
reference_image: str,
|
|
168
|
+
negative_prompt: str,
|
|
169
|
+
width: int,
|
|
170
|
+
height: int,
|
|
171
|
+
steps: int,
|
|
172
|
+
cfg: float,
|
|
173
|
+
seed: int,
|
|
174
|
+
reference_weight: float,
|
|
175
|
+
checkpoint: str
|
|
176
|
+
) -> Dict[str, Any]:
|
|
177
|
+
"""Build IPAdapter workflow for character-consistent generation."""
|
|
178
|
+
|
|
179
|
+
workflow = {
|
|
180
|
+
# Load checkpoint
|
|
181
|
+
"1": {
|
|
182
|
+
"class_type": "CheckpointLoaderSimple",
|
|
183
|
+
"inputs": {
|
|
184
|
+
"ckpt_name": checkpoint
|
|
185
|
+
}
|
|
186
|
+
},
|
|
187
|
+
# Load reference image
|
|
188
|
+
"2": {
|
|
189
|
+
"class_type": "LoadImage",
|
|
190
|
+
"inputs": {
|
|
191
|
+
"image": reference_image
|
|
192
|
+
}
|
|
193
|
+
},
|
|
194
|
+
# Load IPAdapter model
|
|
195
|
+
"3": {
|
|
196
|
+
"class_type": "IPAdapterModelLoader",
|
|
197
|
+
"inputs": {
|
|
198
|
+
"ipadapter_file": "ip-adapter_sd15.safetensors"
|
|
199
|
+
}
|
|
200
|
+
},
|
|
201
|
+
# Load CLIP Vision
|
|
202
|
+
"4": {
|
|
203
|
+
"class_type": "CLIPVisionLoader",
|
|
204
|
+
"inputs": {
|
|
205
|
+
"clip_name": "clip_vision_g.safetensors"
|
|
206
|
+
}
|
|
207
|
+
},
|
|
208
|
+
# Apply IPAdapter
|
|
209
|
+
"5": {
|
|
210
|
+
"class_type": "IPAdapterApply",
|
|
211
|
+
"inputs": {
|
|
212
|
+
"model": ["1", 0],
|
|
213
|
+
"ipadapter": ["3", 0],
|
|
214
|
+
"clip_vision": ["4", 0],
|
|
215
|
+
"image": ["2", 0],
|
|
216
|
+
"weight": reference_weight,
|
|
217
|
+
"noise": 0.0,
|
|
218
|
+
"weight_type": "standard"
|
|
219
|
+
}
|
|
220
|
+
},
|
|
221
|
+
# Positive prompt
|
|
222
|
+
"6": {
|
|
223
|
+
"class_type": "CLIPTextEncode",
|
|
224
|
+
"inputs": {
|
|
225
|
+
"clip": ["1", 1],
|
|
226
|
+
"text": prompt
|
|
227
|
+
}
|
|
228
|
+
},
|
|
229
|
+
# Negative prompt
|
|
230
|
+
"7": {
|
|
231
|
+
"class_type": "CLIPTextEncode",
|
|
232
|
+
"inputs": {
|
|
233
|
+
"clip": ["1", 1],
|
|
234
|
+
"text": negative_prompt
|
|
235
|
+
}
|
|
236
|
+
},
|
|
237
|
+
# Empty latent
|
|
238
|
+
"8": {
|
|
239
|
+
"class_type": "EmptyLatentImage",
|
|
240
|
+
"inputs": {
|
|
241
|
+
"batch_size": 1,
|
|
242
|
+
"height": height,
|
|
243
|
+
"width": width
|
|
244
|
+
}
|
|
245
|
+
},
|
|
246
|
+
# KSampler with IPAdapter-enhanced model
|
|
247
|
+
"9": {
|
|
248
|
+
"class_type": "KSampler",
|
|
249
|
+
"inputs": {
|
|
250
|
+
"cfg": cfg,
|
|
251
|
+
"denoise": 1,
|
|
252
|
+
"latent_image": ["8", 0],
|
|
253
|
+
"model": ["5", 0], # Use IPAdapter-enhanced model
|
|
254
|
+
"negative": ["7", 0],
|
|
255
|
+
"positive": ["6", 0],
|
|
256
|
+
"sampler_name": "euler",
|
|
257
|
+
"scheduler": "normal",
|
|
258
|
+
"seed": seed,
|
|
259
|
+
"steps": steps
|
|
260
|
+
}
|
|
261
|
+
},
|
|
262
|
+
# VAE Decode
|
|
263
|
+
"10": {
|
|
264
|
+
"class_type": "VAEDecode",
|
|
265
|
+
"inputs": {
|
|
266
|
+
"samples": ["9", 0],
|
|
267
|
+
"vae": ["1", 2]
|
|
268
|
+
}
|
|
269
|
+
},
|
|
270
|
+
# Save image
|
|
271
|
+
"11": {
|
|
272
|
+
"class_type": "SaveImage",
|
|
273
|
+
"inputs": {
|
|
274
|
+
"filename_prefix": "vibe_aigc_character",
|
|
275
|
+
"images": ["10", 0]
|
|
276
|
+
}
|
|
277
|
+
}
|
|
278
|
+
}
|
|
279
|
+
|
|
280
|
+
return workflow
|
|
281
|
+
|
|
282
|
+
async def _execute_workflow(self, workflow: Dict[str, Any]) -> GenerationResult:
|
|
283
|
+
"""Execute workflow and wait for completion."""
|
|
284
|
+
payload = {
|
|
285
|
+
"prompt": workflow,
|
|
286
|
+
"client_id": self._client_id
|
|
287
|
+
}
|
|
288
|
+
|
|
289
|
+
try:
|
|
290
|
+
async with aiohttp.ClientSession() as session:
|
|
291
|
+
async with session.post(
|
|
292
|
+
f"{self.config.base_url}/prompt",
|
|
293
|
+
json=payload
|
|
294
|
+
) as resp:
|
|
295
|
+
if resp.status != 200:
|
|
296
|
+
error_text = await resp.text()
|
|
297
|
+
return GenerationResult(
|
|
298
|
+
success=False,
|
|
299
|
+
error=f"Failed to queue prompt: {error_text}"
|
|
300
|
+
)
|
|
301
|
+
result = await resp.json()
|
|
302
|
+
prompt_id = result.get("prompt_id", "")
|
|
303
|
+
|
|
304
|
+
images = await self._wait_for_completion(session, prompt_id)
|
|
305
|
+
|
|
306
|
+
return GenerationResult(
|
|
307
|
+
success=True,
|
|
308
|
+
images=images,
|
|
309
|
+
prompt_id=prompt_id,
|
|
310
|
+
metadata={"type": "character_consistent"}
|
|
311
|
+
)
|
|
312
|
+
|
|
313
|
+
except Exception as e:
|
|
314
|
+
return GenerationResult(
|
|
315
|
+
success=False,
|
|
316
|
+
error=str(e)
|
|
317
|
+
)
|
|
318
|
+
|
|
319
|
+
async def _wait_for_completion(
|
|
320
|
+
self,
|
|
321
|
+
session: aiohttp.ClientSession,
|
|
322
|
+
prompt_id: str,
|
|
323
|
+
timeout: float = 300,
|
|
324
|
+
poll_interval: float = 0.5
|
|
325
|
+
) -> List[str]:
|
|
326
|
+
"""Wait for generation to complete."""
|
|
327
|
+
import time
|
|
328
|
+
start_time = time.time()
|
|
329
|
+
|
|
330
|
+
while time.time() - start_time < timeout:
|
|
331
|
+
async with session.get(f"{self.config.base_url}/history/{prompt_id}") as resp:
|
|
332
|
+
if resp.status == 200:
|
|
333
|
+
history = await resp.json()
|
|
334
|
+
|
|
335
|
+
if prompt_id in history:
|
|
336
|
+
prompt_data = history[prompt_id]
|
|
337
|
+
|
|
338
|
+
if prompt_data.get("status", {}).get("completed", False):
|
|
339
|
+
images = []
|
|
340
|
+
outputs = prompt_data.get("outputs", {})
|
|
341
|
+
|
|
342
|
+
for node_id, node_output in outputs.items():
|
|
343
|
+
if "images" in node_output:
|
|
344
|
+
for img in node_output["images"]:
|
|
345
|
+
filename = img.get("filename")
|
|
346
|
+
subfolder = img.get("subfolder", "")
|
|
347
|
+
if filename:
|
|
348
|
+
url = f"{self.config.base_url}/view?filename={filename}"
|
|
349
|
+
if subfolder:
|
|
350
|
+
url += f"&subfolder={subfolder}"
|
|
351
|
+
images.append(url)
|
|
352
|
+
|
|
353
|
+
return images
|
|
354
|
+
|
|
355
|
+
await asyncio.sleep(poll_interval)
|
|
356
|
+
|
|
357
|
+
raise TimeoutError(f"Generation timed out after {timeout}s")
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
class CharacterConsistentTool(BaseTool):
|
|
361
|
+
"""Tool for generating character-consistent images."""
|
|
362
|
+
|
|
363
|
+
def __init__(self, config: Optional[ComfyUIConfig] = None):
|
|
364
|
+
self.backend = IPAdapterBackend(config)
|
|
365
|
+
self.character_bank = CharacterBank()
|
|
366
|
+
self._spec = ToolSpec(
|
|
367
|
+
name="character_consistent",
|
|
368
|
+
description="Generate images with character consistency using IPAdapter",
|
|
369
|
+
category=ToolCategory.IMAGE,
|
|
370
|
+
input_schema={
|
|
371
|
+
"type": "object",
|
|
372
|
+
"required": ["prompt"],
|
|
373
|
+
"properties": {
|
|
374
|
+
"prompt": {"type": "string"},
|
|
375
|
+
"character_name": {"type": "string", "description": "Name of character from bank"},
|
|
376
|
+
"reference_image": {"type": "string", "description": "Direct path to reference image"},
|
|
377
|
+
"reference_weight": {"type": "number", "default": 0.8},
|
|
378
|
+
"negative_prompt": {"type": "string"},
|
|
379
|
+
"width": {"type": "integer", "default": 512},
|
|
380
|
+
"height": {"type": "integer", "default": 512},
|
|
381
|
+
"steps": {"type": "integer", "default": 20},
|
|
382
|
+
"cfg": {"type": "number", "default": 7.0}
|
|
383
|
+
}
|
|
384
|
+
},
|
|
385
|
+
output_schema={
|
|
386
|
+
"type": "object",
|
|
387
|
+
"properties": {
|
|
388
|
+
"images": {"type": "array", "items": {"type": "string"}},
|
|
389
|
+
"prompt_id": {"type": "string"}
|
|
390
|
+
}
|
|
391
|
+
}
|
|
392
|
+
)
|
|
393
|
+
|
|
394
|
+
@property
|
|
395
|
+
def spec(self) -> ToolSpec:
|
|
396
|
+
return self._spec
|
|
397
|
+
|
|
398
|
+
def add_character(self, character: CharacterReference) -> None:
|
|
399
|
+
"""Add a character to the bank."""
|
|
400
|
+
self.character_bank.add(character)
|
|
401
|
+
|
|
402
|
+
async def execute(
|
|
403
|
+
self,
|
|
404
|
+
inputs: Dict[str, Any],
|
|
405
|
+
context: Optional[Dict[str, Any]] = None
|
|
406
|
+
) -> ToolResult:
|
|
407
|
+
"""Execute character-consistent generation."""
|
|
408
|
+
prompt = inputs.get("prompt", "")
|
|
409
|
+
if not prompt:
|
|
410
|
+
return ToolResult(success=False, output=None, error="No prompt provided")
|
|
411
|
+
|
|
412
|
+
# Get reference image
|
|
413
|
+
reference_image = inputs.get("reference_image")
|
|
414
|
+
character_name = inputs.get("character_name")
|
|
415
|
+
reference_weight = inputs.get("reference_weight", 0.8)
|
|
416
|
+
|
|
417
|
+
if character_name:
|
|
418
|
+
char = self.character_bank.get(character_name)
|
|
419
|
+
if char:
|
|
420
|
+
reference_image = char.reference_image
|
|
421
|
+
reference_weight = char.weight
|
|
422
|
+
else:
|
|
423
|
+
return ToolResult(
|
|
424
|
+
success=False,
|
|
425
|
+
output=None,
|
|
426
|
+
error=f"Character '{character_name}' not found in bank"
|
|
427
|
+
)
|
|
428
|
+
|
|
429
|
+
if not reference_image:
|
|
430
|
+
return ToolResult(
|
|
431
|
+
success=False,
|
|
432
|
+
output=None,
|
|
433
|
+
error="No reference_image or character_name provided"
|
|
434
|
+
)
|
|
435
|
+
|
|
436
|
+
result = await self.backend.generate_with_reference(
|
|
437
|
+
prompt=prompt,
|
|
438
|
+
reference_image=reference_image,
|
|
439
|
+
negative_prompt=inputs.get("negative_prompt", ""),
|
|
440
|
+
width=inputs.get("width", 512),
|
|
441
|
+
height=inputs.get("height", 512),
|
|
442
|
+
steps=inputs.get("steps", 20),
|
|
443
|
+
cfg=inputs.get("cfg", 7.0),
|
|
444
|
+
reference_weight=reference_weight
|
|
445
|
+
)
|
|
446
|
+
|
|
447
|
+
if result.success:
|
|
448
|
+
return ToolResult(
|
|
449
|
+
success=True,
|
|
450
|
+
output={
|
|
451
|
+
"images": result.images,
|
|
452
|
+
"prompt_id": result.prompt_id
|
|
453
|
+
},
|
|
454
|
+
metadata=result.metadata
|
|
455
|
+
)
|
|
456
|
+
else:
|
|
457
|
+
return ToolResult(success=False, output=None, error=result.error)
|
|
@@ -0,0 +1,388 @@
|
|
|
1
|
+
"""Video generation backend using AnimateDiff.
|
|
2
|
+
|
|
3
|
+
This extends vibe-aigc from images to video, completing the
|
|
4
|
+
multimodal content generation capabilities.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import asyncio
|
|
8
|
+
import json
|
|
9
|
+
import uuid
|
|
10
|
+
import aiohttp
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
from typing import Any, Dict, List, Optional
|
|
13
|
+
from dataclasses import dataclass, field
|
|
14
|
+
|
|
15
|
+
from .comfyui import ComfyUIConfig, GenerationResult
|
|
16
|
+
from .tools import BaseTool, ToolResult, ToolSpec, ToolCategory
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class VideoConfig:
|
|
21
|
+
"""Configuration for video generation."""
|
|
22
|
+
frames: int = 16 # Number of frames (16, 24, or 32 typical)
|
|
23
|
+
fps: int = 8 # Frames per second
|
|
24
|
+
motion_scale: float = 1.0 # Motion intensity (0.5-1.5)
|
|
25
|
+
loop: bool = False # Whether to make it loop
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class AnimateDiffBackend:
|
|
29
|
+
"""Backend for video generation via AnimateDiff in ComfyUI."""
|
|
30
|
+
|
|
31
|
+
def __init__(self, config: Optional[ComfyUIConfig] = None):
|
|
32
|
+
self.config = config or ComfyUIConfig()
|
|
33
|
+
self._client_id = str(uuid.uuid4())
|
|
34
|
+
|
|
35
|
+
async def is_available(self) -> bool:
|
|
36
|
+
"""Check if ComfyUI is running with AnimateDiff nodes."""
|
|
37
|
+
try:
|
|
38
|
+
async with aiohttp.ClientSession() as session:
|
|
39
|
+
async with session.get(
|
|
40
|
+
f"{self.config.base_url}/object_info",
|
|
41
|
+
timeout=aiohttp.ClientTimeout(total=5)
|
|
42
|
+
) as resp:
|
|
43
|
+
if resp.status != 200:
|
|
44
|
+
return False
|
|
45
|
+
obj_info = await resp.json()
|
|
46
|
+
# Check for AnimateDiff nodes
|
|
47
|
+
return "ADE_AnimateDiffLoaderGen1" in obj_info or "AnimateDiffLoaderV1" in obj_info
|
|
48
|
+
except Exception:
|
|
49
|
+
return False
|
|
50
|
+
|
|
51
|
+
async def generate_video(
|
|
52
|
+
self,
|
|
53
|
+
prompt: str,
|
|
54
|
+
negative_prompt: str = "",
|
|
55
|
+
width: int = 512,
|
|
56
|
+
height: int = 512,
|
|
57
|
+
frames: int = 16,
|
|
58
|
+
fps: int = 8,
|
|
59
|
+
steps: int = 20,
|
|
60
|
+
cfg: float = 7.0,
|
|
61
|
+
seed: Optional[int] = None,
|
|
62
|
+
motion_model: str = "v3_sd15_mm.ckpt",
|
|
63
|
+
checkpoint: str = "v1-5-pruned-emaonly.safetensors"
|
|
64
|
+
) -> GenerationResult:
|
|
65
|
+
"""Generate a video using AnimateDiff.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
prompt: Text description of the video
|
|
69
|
+
negative_prompt: What to avoid
|
|
70
|
+
width: Video width
|
|
71
|
+
height: Video height
|
|
72
|
+
frames: Number of frames (16, 24, 32)
|
|
73
|
+
fps: Frames per second for output
|
|
74
|
+
steps: Sampling steps
|
|
75
|
+
cfg: Guidance scale
|
|
76
|
+
seed: Random seed
|
|
77
|
+
motion_model: AnimateDiff motion model
|
|
78
|
+
checkpoint: SD checkpoint to use
|
|
79
|
+
|
|
80
|
+
Returns:
|
|
81
|
+
GenerationResult with video file paths
|
|
82
|
+
"""
|
|
83
|
+
if seed is None:
|
|
84
|
+
import random
|
|
85
|
+
seed = random.randint(0, 2**32 - 1)
|
|
86
|
+
|
|
87
|
+
workflow = self._build_animatediff_workflow(
|
|
88
|
+
prompt=prompt,
|
|
89
|
+
negative_prompt=negative_prompt,
|
|
90
|
+
width=width,
|
|
91
|
+
height=height,
|
|
92
|
+
frames=frames,
|
|
93
|
+
steps=steps,
|
|
94
|
+
cfg=cfg,
|
|
95
|
+
seed=seed,
|
|
96
|
+
motion_model=motion_model,
|
|
97
|
+
checkpoint=checkpoint,
|
|
98
|
+
fps=fps
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
return await self._execute_workflow(workflow)
|
|
102
|
+
|
|
103
|
+
def _build_animatediff_workflow(
|
|
104
|
+
self,
|
|
105
|
+
prompt: str,
|
|
106
|
+
negative_prompt: str,
|
|
107
|
+
width: int,
|
|
108
|
+
height: int,
|
|
109
|
+
frames: int,
|
|
110
|
+
steps: int,
|
|
111
|
+
cfg: float,
|
|
112
|
+
seed: int,
|
|
113
|
+
motion_model: str,
|
|
114
|
+
checkpoint: str,
|
|
115
|
+
fps: int
|
|
116
|
+
) -> Dict[str, Any]:
|
|
117
|
+
"""Build AnimateDiff workflow in ComfyUI API format.
|
|
118
|
+
|
|
119
|
+
Uses ADE_UseEvolvedSampling for proper AnimateDiff integration.
|
|
120
|
+
"""
|
|
121
|
+
|
|
122
|
+
workflow = {
|
|
123
|
+
# Load checkpoint
|
|
124
|
+
"1": {
|
|
125
|
+
"class_type": "CheckpointLoaderSimple",
|
|
126
|
+
"inputs": {
|
|
127
|
+
"ckpt_name": checkpoint
|
|
128
|
+
}
|
|
129
|
+
},
|
|
130
|
+
# Load AnimateDiff motion model
|
|
131
|
+
"2": {
|
|
132
|
+
"class_type": "ADE_LoadAnimateDiffModel",
|
|
133
|
+
"inputs": {
|
|
134
|
+
"model_name": motion_model
|
|
135
|
+
}
|
|
136
|
+
},
|
|
137
|
+
# Convert motion model to M_MODELS type
|
|
138
|
+
"2b": {
|
|
139
|
+
"class_type": "ADE_ApplyAnimateDiffModelSimple",
|
|
140
|
+
"inputs": {
|
|
141
|
+
"motion_model": ["2", 0]
|
|
142
|
+
}
|
|
143
|
+
},
|
|
144
|
+
# Apply AnimateDiff with evolved sampling
|
|
145
|
+
"3": {
|
|
146
|
+
"class_type": "ADE_UseEvolvedSampling",
|
|
147
|
+
"inputs": {
|
|
148
|
+
"model": ["1", 0],
|
|
149
|
+
"m_models": ["2b", 0],
|
|
150
|
+
"beta_schedule": "sqrt_linear (AnimateDiff)"
|
|
151
|
+
}
|
|
152
|
+
},
|
|
153
|
+
# Positive prompt
|
|
154
|
+
"4": {
|
|
155
|
+
"class_type": "CLIPTextEncode",
|
|
156
|
+
"inputs": {
|
|
157
|
+
"clip": ["1", 1],
|
|
158
|
+
"text": prompt
|
|
159
|
+
}
|
|
160
|
+
},
|
|
161
|
+
# Negative prompt
|
|
162
|
+
"5": {
|
|
163
|
+
"class_type": "CLIPTextEncode",
|
|
164
|
+
"inputs": {
|
|
165
|
+
"clip": ["1", 1],
|
|
166
|
+
"text": negative_prompt
|
|
167
|
+
}
|
|
168
|
+
},
|
|
169
|
+
# Empty latent batch for video frames
|
|
170
|
+
"6": {
|
|
171
|
+
"class_type": "EmptyLatentImage",
|
|
172
|
+
"inputs": {
|
|
173
|
+
"batch_size": frames,
|
|
174
|
+
"height": height,
|
|
175
|
+
"width": width
|
|
176
|
+
}
|
|
177
|
+
},
|
|
178
|
+
# KSampler with AnimateDiff model
|
|
179
|
+
"7": {
|
|
180
|
+
"class_type": "KSampler",
|
|
181
|
+
"inputs": {
|
|
182
|
+
"cfg": cfg,
|
|
183
|
+
"denoise": 1,
|
|
184
|
+
"latent_image": ["6", 0],
|
|
185
|
+
"model": ["3", 0],
|
|
186
|
+
"negative": ["5", 0],
|
|
187
|
+
"positive": ["4", 0],
|
|
188
|
+
"sampler_name": "euler_ancestral",
|
|
189
|
+
"scheduler": "normal",
|
|
190
|
+
"seed": seed,
|
|
191
|
+
"steps": steps
|
|
192
|
+
}
|
|
193
|
+
},
|
|
194
|
+
# VAE Decode
|
|
195
|
+
"8": {
|
|
196
|
+
"class_type": "VAEDecode",
|
|
197
|
+
"inputs": {
|
|
198
|
+
"samples": ["7", 0],
|
|
199
|
+
"vae": ["1", 2]
|
|
200
|
+
}
|
|
201
|
+
},
|
|
202
|
+
# Use AnimateDiff's built-in combiner for GIF output
|
|
203
|
+
"9": {
|
|
204
|
+
"class_type": "ADE_AnimateDiffCombine",
|
|
205
|
+
"inputs": {
|
|
206
|
+
"images": ["8", 0],
|
|
207
|
+
"frame_rate": fps,
|
|
208
|
+
"loop_count": 0,
|
|
209
|
+
"format": "image/gif",
|
|
210
|
+
"pingpong": False,
|
|
211
|
+
"save_image": True,
|
|
212
|
+
"filename_prefix": "vibe_aigc_video"
|
|
213
|
+
}
|
|
214
|
+
}
|
|
215
|
+
}
|
|
216
|
+
|
|
217
|
+
return workflow
|
|
218
|
+
|
|
219
|
+
async def _execute_workflow(self, workflow: Dict[str, Any]) -> GenerationResult:
|
|
220
|
+
"""Execute workflow and wait for completion."""
|
|
221
|
+
payload = {
|
|
222
|
+
"prompt": workflow,
|
|
223
|
+
"client_id": self._client_id
|
|
224
|
+
}
|
|
225
|
+
|
|
226
|
+
try:
|
|
227
|
+
async with aiohttp.ClientSession() as session:
|
|
228
|
+
# Queue the prompt
|
|
229
|
+
async with session.post(
|
|
230
|
+
f"{self.config.base_url}/prompt",
|
|
231
|
+
json=payload
|
|
232
|
+
) as resp:
|
|
233
|
+
if resp.status != 200:
|
|
234
|
+
error_text = await resp.text()
|
|
235
|
+
return GenerationResult(
|
|
236
|
+
success=False,
|
|
237
|
+
error=f"Failed to queue prompt: {error_text}"
|
|
238
|
+
)
|
|
239
|
+
result = await resp.json()
|
|
240
|
+
prompt_id = result.get("prompt_id", "")
|
|
241
|
+
|
|
242
|
+
# Wait for completion
|
|
243
|
+
videos = await self._wait_for_completion(session, prompt_id)
|
|
244
|
+
|
|
245
|
+
return GenerationResult(
|
|
246
|
+
success=True,
|
|
247
|
+
images=videos, # reusing images field for video paths
|
|
248
|
+
prompt_id=prompt_id,
|
|
249
|
+
metadata={"type": "video", "workflow": "animatediff"}
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
except Exception as e:
|
|
253
|
+
return GenerationResult(
|
|
254
|
+
success=False,
|
|
255
|
+
error=str(e)
|
|
256
|
+
)
|
|
257
|
+
|
|
258
|
+
async def _wait_for_completion(
|
|
259
|
+
self,
|
|
260
|
+
session: aiohttp.ClientSession,
|
|
261
|
+
prompt_id: str,
|
|
262
|
+
timeout: float = 600, # Videos take longer
|
|
263
|
+
poll_interval: float = 1.0
|
|
264
|
+
) -> List[str]:
|
|
265
|
+
"""Wait for video generation to complete."""
|
|
266
|
+
import time
|
|
267
|
+
start_time = time.time()
|
|
268
|
+
|
|
269
|
+
while time.time() - start_time < timeout:
|
|
270
|
+
async with session.get(f"{self.config.base_url}/history/{prompt_id}") as resp:
|
|
271
|
+
if resp.status == 200:
|
|
272
|
+
history = await resp.json()
|
|
273
|
+
|
|
274
|
+
if prompt_id in history:
|
|
275
|
+
prompt_data = history[prompt_id]
|
|
276
|
+
|
|
277
|
+
if prompt_data.get("status", {}).get("completed", False):
|
|
278
|
+
videos = []
|
|
279
|
+
outputs = prompt_data.get("outputs", {})
|
|
280
|
+
|
|
281
|
+
for node_id, node_output in outputs.items():
|
|
282
|
+
# Check for animated outputs (webp, gif)
|
|
283
|
+
if "images" in node_output:
|
|
284
|
+
for img in node_output["images"]:
|
|
285
|
+
filename = img.get("filename")
|
|
286
|
+
subfolder = img.get("subfolder", "")
|
|
287
|
+
if filename:
|
|
288
|
+
# Animated webp/gif files
|
|
289
|
+
url = f"{self.config.base_url}/view?filename={filename}"
|
|
290
|
+
if subfolder:
|
|
291
|
+
url += f"&subfolder={subfolder}"
|
|
292
|
+
videos.append(url)
|
|
293
|
+
# Also check gifs field (some nodes use this)
|
|
294
|
+
if "gifs" in node_output:
|
|
295
|
+
for gif in node_output["gifs"]:
|
|
296
|
+
filename = gif.get("filename")
|
|
297
|
+
subfolder = gif.get("subfolder", "")
|
|
298
|
+
if filename:
|
|
299
|
+
url = f"{self.config.base_url}/view?filename={filename}"
|
|
300
|
+
if subfolder:
|
|
301
|
+
url += f"&subfolder={subfolder}"
|
|
302
|
+
videos.append(url)
|
|
303
|
+
|
|
304
|
+
return videos
|
|
305
|
+
|
|
306
|
+
await asyncio.sleep(poll_interval)
|
|
307
|
+
|
|
308
|
+
raise TimeoutError(f"Video generation timed out after {timeout}s")
|
|
309
|
+
|
|
310
|
+
|
|
311
|
+
class AnimateDiffTool(BaseTool):
|
|
312
|
+
"""Tool for generating videos via AnimateDiff."""
|
|
313
|
+
|
|
314
|
+
def __init__(self, config: Optional[ComfyUIConfig] = None):
|
|
315
|
+
self.backend = AnimateDiffBackend(config)
|
|
316
|
+
self._spec = ToolSpec(
|
|
317
|
+
name="animatediff_video",
|
|
318
|
+
description="Generate short video clips using AnimateDiff (local)",
|
|
319
|
+
category=ToolCategory.VIDEO,
|
|
320
|
+
input_schema={
|
|
321
|
+
"type": "object",
|
|
322
|
+
"required": ["prompt"],
|
|
323
|
+
"properties": {
|
|
324
|
+
"prompt": {"type": "string", "description": "Video description"},
|
|
325
|
+
"negative_prompt": {"type": "string"},
|
|
326
|
+
"width": {"type": "integer", "default": 512},
|
|
327
|
+
"height": {"type": "integer", "default": 512},
|
|
328
|
+
"frames": {"type": "integer", "default": 16, "description": "Number of frames (16, 24, 32)"},
|
|
329
|
+
"fps": {"type": "integer", "default": 8},
|
|
330
|
+
"steps": {"type": "integer", "default": 20},
|
|
331
|
+
"cfg": {"type": "number", "default": 7.0},
|
|
332
|
+
"seed": {"type": "integer"}
|
|
333
|
+
}
|
|
334
|
+
},
|
|
335
|
+
output_schema={
|
|
336
|
+
"type": "object",
|
|
337
|
+
"properties": {
|
|
338
|
+
"videos": {"type": "array", "items": {"type": "string"}},
|
|
339
|
+
"prompt_id": {"type": "string"}
|
|
340
|
+
}
|
|
341
|
+
}
|
|
342
|
+
)
|
|
343
|
+
|
|
344
|
+
@property
|
|
345
|
+
def spec(self) -> ToolSpec:
|
|
346
|
+
return self._spec
|
|
347
|
+
|
|
348
|
+
async def execute(
|
|
349
|
+
self,
|
|
350
|
+
inputs: Dict[str, Any],
|
|
351
|
+
context: Optional[Dict[str, Any]] = None
|
|
352
|
+
) -> ToolResult:
|
|
353
|
+
"""Execute video generation."""
|
|
354
|
+
prompt = inputs.get("prompt", "")
|
|
355
|
+
if not prompt:
|
|
356
|
+
return ToolResult(success=False, output=None, error="No prompt provided")
|
|
357
|
+
|
|
358
|
+
# Check availability
|
|
359
|
+
if not await self.backend.is_available():
|
|
360
|
+
return ToolResult(
|
|
361
|
+
success=False,
|
|
362
|
+
output=None,
|
|
363
|
+
error="AnimateDiff not available. Install ComfyUI-AnimateDiff-Evolved node."
|
|
364
|
+
)
|
|
365
|
+
|
|
366
|
+
result = await self.backend.generate_video(
|
|
367
|
+
prompt=prompt,
|
|
368
|
+
negative_prompt=inputs.get("negative_prompt", ""),
|
|
369
|
+
width=inputs.get("width", 512),
|
|
370
|
+
height=inputs.get("height", 512),
|
|
371
|
+
frames=inputs.get("frames", 16),
|
|
372
|
+
fps=inputs.get("fps", 8),
|
|
373
|
+
steps=inputs.get("steps", 20),
|
|
374
|
+
cfg=inputs.get("cfg", 7.0),
|
|
375
|
+
seed=inputs.get("seed")
|
|
376
|
+
)
|
|
377
|
+
|
|
378
|
+
if result.success:
|
|
379
|
+
return ToolResult(
|
|
380
|
+
success=True,
|
|
381
|
+
output={
|
|
382
|
+
"videos": result.images,
|
|
383
|
+
"prompt_id": result.prompt_id
|
|
384
|
+
},
|
|
385
|
+
metadata=result.metadata
|
|
386
|
+
)
|
|
387
|
+
else:
|
|
388
|
+
return ToolResult(success=False, output=None, error=result.error)
|
|
@@ -24,6 +24,7 @@ tests/test_workflow_resume.py
|
|
|
24
24
|
vibe_aigc/__init__.py
|
|
25
25
|
vibe_aigc/agents.py
|
|
26
26
|
vibe_aigc/assets.py
|
|
27
|
+
vibe_aigc/character.py
|
|
27
28
|
vibe_aigc/cli.py
|
|
28
29
|
vibe_aigc/comfyui.py
|
|
29
30
|
vibe_aigc/executor.py
|
|
@@ -34,6 +35,7 @@ vibe_aigc/persistence.py
|
|
|
34
35
|
vibe_aigc/planner.py
|
|
35
36
|
vibe_aigc/tools.py
|
|
36
37
|
vibe_aigc/tools_multimodal.py
|
|
38
|
+
vibe_aigc/video.py
|
|
37
39
|
vibe_aigc/visualization.py
|
|
38
40
|
vibe_aigc.egg-info/PKG-INFO
|
|
39
41
|
vibe_aigc.egg-info/SOURCES.txt
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|