25e57d8221
- Replace placeholder workflow with complete SD1.5 pipeline - Support dynamic prompt, negative_prompt, steps, seed, width, height - Nodes: CheckpointLoader -> CLIP -> KSampler -> VAE -> SaveImage Co-Authored-By: Warp <agent@warp.dev>
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
# services/comfy-agent/app/api.py
|
|
from fastapi import APIRouter, HTTPException
|
|
from .models import GenerateImageRequest, GenerateVideoRequest, JobStatus
|
|
from .jobs import JOB_STORE
|
|
from .worker import enqueue
|
|
|
|
router = APIRouter()
|
|
|
|
def _build_workflow_t2i(req: GenerateImageRequest) -> dict:
|
|
# Basic SD 1.5 workflow
|
|
# Node structure: CheckpointLoader -> CLIP Encode -> KSampler -> VAE Decode -> SaveImage
|
|
return {
|
|
"3": {
|
|
"inputs": {
|
|
"seed": req.seed if req.seed else 42,
|
|
"steps": req.steps,
|
|
"cfg": 7.0,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1,
|
|
"model": ["4", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
"latent_image": ["5", 0]
|
|
},
|
|
"class_type": "KSampler"
|
|
},
|
|
"4": {
|
|
"inputs": {
|
|
"ckpt_name": "v1-5-pruned-emaonly.safetensors"
|
|
},
|
|
"class_type": "CheckpointLoaderSimple"
|
|
},
|
|
"5": {
|
|
"inputs": {
|
|
"width": req.width,
|
|
"height": req.height,
|
|
"batch_size": 1
|
|
},
|
|
"class_type": "EmptyLatentImage"
|
|
},
|
|
"6": {
|
|
"inputs": {
|
|
"text": req.prompt,
|
|
"clip": ["4", 1]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"7": {
|
|
"inputs": {
|
|
"text": req.negative_prompt if req.negative_prompt else "text, watermark, blurry",
|
|
"clip": ["4", 1]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"8": {
|
|
"inputs": {
|
|
"samples": ["3", 0],
|
|
"vae": ["4", 2]
|
|
},
|
|
"class_type": "VAEDecode"
|
|
},
|
|
"9": {
|
|
"inputs": {
|
|
"filename_prefix": "comfy-agent",
|
|
"images": ["8", 0]
|
|
},
|
|
"class_type": "SaveImage"
|
|
}
|
|
}
|
|
|
|
def _build_workflow_t2v(req: GenerateVideoRequest) -> dict:
|
|
# MVP placeholder for LTX-2 pipeline; replace with actual LTX-2 workflow.
|
|
return {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": req.prompt, "clip": ["2", 0]}},
|
|
# TODO: Add complete workflow JSON for text-to-video with LTX-2
|
|
}
|
|
|
|
@router.post("/generate/image", response_model=JobStatus)
|
|
async def generate_image(req: GenerateImageRequest):
|
|
job = JOB_STORE.create("text-to-image")
|
|
graph = _build_workflow_t2i(req)
|
|
enqueue(job.job_id, "text-to-image", graph)
|
|
return JOB_STORE.get(job.job_id)
|
|
|
|
@router.post("/generate/video", response_model=JobStatus)
|
|
async def generate_video(req: GenerateVideoRequest):
|
|
job = JOB_STORE.create("text-to-video")
|
|
graph = _build_workflow_t2v(req)
|
|
enqueue(job.job_id, "text-to-video", graph)
|
|
return JOB_STORE.get(job.job_id)
|
|
|
|
@router.get("/status/{job_id}", response_model=JobStatus)
|
|
async def status(job_id: str):
|
|
job = JOB_STORE.get(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job_not_found")
|
|
return job
|
|
|
|
@router.get("/result/{job_id}", response_model=JobStatus)
|
|
async def result(job_id: str):
|
|
job = JOB_STORE.get(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job_not_found")
|
|
return job
|