Skip to content

Commit 8bf3d8e

Browse files
eliteproxeliteprox
committed
use resolution from prompt
1 parent 321290b commit 8bf3d8e

File tree

6 files changed

+66
-146
lines changed

6 files changed

+66
-146
lines changed

‎runner/app/live/pipelines/comfyui.py‎

Lines changed: 8 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -11,53 +11,20 @@
1111
from .interface import Pipeline
1212
from comfystream.client import ComfyStreamClient
1313
from trickle import VideoFrame, VideoOutput
14+
from utils import ComfyUtils
1415

1516
import logging
1617

1718
COMFY_UI_WORKSPACE_ENV = "COMFY_UI_WORKSPACE"
1819
WARMUP_RUNS = 1
19-
class ComfyUtils:
20-
DEFAULT_WIDTH = 384
21-
DEFAULT_HEIGHT = 704
22-
@staticmethod
23-
def get_default_workflow_json():
24-
_default_workflow_path = pathlib.Path(__file__).parent.absolute() / "comfyui_default_workflow.json"
25-
with open(_default_workflow_path, 'r') as f:
26-
return json.load(f)
27-
28-
@staticmethod
29-
def get_latent_image_dimensions(workflow: dict | None) -> tuple[int, int]:
30-
"""Get dimensions from the EmptyLatentImage node in the workflow.
31-
32-
Args:
33-
workflow: The workflow JSON dictionary
34-
35-
Returns:
36-
Tuple of (width, height) from the latent image. Returns default dimensions if not found or on error.
37-
"""
38-
39-
if workflow is None:
40-
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
41-
42-
try:
43-
for node_id, node in workflow.items():
44-
if node.get("class_type") == "EmptyLatentImage":
45-
inputs = node.get("inputs", {})
46-
width = inputs.get("width")
47-
height = inputs.get("height")
48-
if width is not None and height is not None:
49-
return width, height
50-
logging.warning("Incomplete dimensions in latent image node")
51-
break
52-
except Exception as e:
53-
logging.warning(f"Failed to extract dimensions from workflow: {e}")
54-
55-
# Return defaults if dimensions not found or on any error
56-
logging.info(f"Using default dimensions {ComfyUtils.DEFAULT_WIDTH}x{ComfyUtils.DEFAULT_HEIGHT}")
57-
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
20+
21+
def get_default_workflow_json():
22+
_default_workflow_path = pathlib.Path(__file__).parent.absolute() / "comfyui_default_workflow.json"
23+
with open(_default_workflow_path, 'r') as f:
24+
return json.load(f)
5825

5926
# Get the default workflow json during startup
60-
DEFAULT_WORKFLOW_JSON = ComfyUtils.get_default_workflow_json()
27+
DEFAULT_WORKFLOW_JSON = get_default_workflow_json()
6128

6229
class ComfyUIParams(BaseModel):
6330
class Config:
@@ -149,30 +116,4 @@ async def update_params(self, **params):
149116
async def stop(self):
150117
logging.info("Stopping ComfyUI pipeline")
151118
await self.client.cleanup()
152-
logging.info("ComfyUI pipeline stopped")
153-
154-
# async def get_latent_image_dimensions(self, workflow: dict) -> tuple[int, int] | None:
155-
# """Get dimensions from the EmptyLatentImage node in the workflow.
156-
157-
# Args:
158-
# workflow: The workflow JSON dictionary
159-
160-
# Returns:
161-
# Tuple of (width, height) from the latent image. Returns default dimensions if not found or on error.
162-
# """
163-
# try:
164-
# for node_id, node in workflow.items():
165-
# if node.get("class_type") == "EmptyLatentImage":
166-
# inputs = node.get("inputs", {})
167-
# width = inputs.get("width")
168-
# height = inputs.get("height")
169-
# if width is not None and height is not None:
170-
# return width, height
171-
# logging.warning("Incomplete dimensions in latent image node")
172-
# break
173-
# except Exception as e:
174-
# logging.warning(f"Failed to extract dimensions from workflow: {e}")
175-
176-
# # Return defaults if dimensions not found or on any error
177-
# logging.info(f"Using default dimensions {DEFAULT_WIDTH}x{DEFAULT_HEIGHT}")
178-
# return DEFAULT_WIDTH, DEFAULT_HEIGHT
119+
logging.info("ComfyUI pipeline stopped")

‎runner/app/live/streamer/process.py‎

Lines changed: 4 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from pipelines import load_pipeline, Pipeline
1212
from log import config_logging, config_logging_fields, log_timing
1313
from trickle import InputFrame, AudioFrame, VideoFrame, OutputFrame, VideoOutput, AudioOutput
14+
from utils import ComfyUtils
1415

1516

1617
class PipelineProcess:
@@ -165,20 +166,9 @@ async def _initialize_pipeline(self):
165166

166167
with log_timing(f"PipelineProcess: Pipeline loading with {params}"):
167168
pipeline = load_pipeline(self.pipeline_name)
168-
if params.get('prompt') is None:
169-
# Parse resolution from the default workflow
170-
width, height = ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
171-
else:
172-
# Parse resolution from the provided workflow
173-
prompt = params.get('prompt')
174-
if prompt is type(dict):
175-
width, height = ComfyUtils.get_latent_image_dimensions(prompt)
176-
177-
# TODO pass through to streamer
178-
# params.update({
179-
# 'width': width,
180-
# 'height': height
181-
# })
169+
170+
# TODO: We may need to call reset_stream when resolution is changed and start the pipeline again
171+
# Changing the engine causes issues, maybe cleanup related
182172
await pipeline.initialize(**params)
183173
return pipeline
184174
except Exception as e:
@@ -348,38 +338,3 @@ def clear_queue(queue):
348338
queue.get_nowait() # Remove items without blocking
349339
except Exception as e:
350340
logging.error(f"Error while clearing queue: {e}")
351-
352-
class ComfyUtils:
353-
DEFAULT_WIDTH = 384
354-
DEFAULT_HEIGHT = 704
355-
356-
@staticmethod
357-
def get_latent_image_dimensions(workflow: dict | None) -> tuple[int, int]:
358-
"""Get dimensions from the EmptyLatentImage node in the workflow.
359-
360-
Args:
361-
workflow: The workflow JSON dictionary
362-
363-
Returns:
364-
Tuple of (width, height) from the latent image. Returns default dimensions if not found or on error.
365-
"""
366-
367-
if workflow is None:
368-
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
369-
370-
try:
371-
for node_id, node in workflow.items():
372-
if node.get("class_type") == "EmptyLatentImage":
373-
inputs = node.get("inputs", {})
374-
width = inputs.get("width")
375-
height = inputs.get("height")
376-
if width is not None and height is not None:
377-
return width, height
378-
logging.warning("Incomplete dimensions in latent image node")
379-
break
380-
except Exception as e:
381-
logging.warning(f"Failed to extract dimensions from workflow: {e}")
382-
383-
# Return defaults if dimensions not found or on any error
384-
logging.info(f"Using default dimensions {ComfyUtils.DEFAULT_WIDTH}x{ComfyUtils.DEFAULT_HEIGHT}")
385-
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT

‎runner/app/live/streamer/protocol/trickle.py‎

Lines changed: 4 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from PIL import Image
88

99
from trickle import media, TricklePublisher, TrickleSubscriber, InputFrame, OutputFrame, AudioFrame, AudioOutput
10-
10+
from utils import ComfyUtils
1111
from .protocol import StreamProtocol
1212
from .last_value_cache import LastValueCache
1313

@@ -23,8 +23,8 @@ def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[s
2323
self.events_publisher = None
2424
self.subscribe_task = None
2525
self.publish_task = None
26-
self.width = 384 # Default values
27-
self.height = 704
26+
self.width = ComfyUtils.DEFAULT_WIDTH
27+
self.height = ComfyUtils.DEFAULT_HEIGHT
2828

2929
async def start(self):
3030
self.subscribe_queue = queue.Queue[InputFrame]()
@@ -126,26 +126,7 @@ async def control_loop(self, done: asyncio.Event) -> AsyncGenerator[dict, None]:
126126
# Ignore periodic keepalive messages
127127
continue
128128

129-
# Handle resolution changes
130-
# TODO: This should be on the input (encode, subscribe), not output
131-
# if 'width' in data or 'height' in data:
132-
# new_width = data.get('width', self.width)
133-
# new_height = data.get('height', self.height)
134-
# if new_width != self.width or new_height != self.height:
135-
# logging.info(f"Updating resolution from {self.width}x{self.height} to {new_width}x{new_height}")
136-
# self.width = new_width
137-
# self.height = new_height
138-
# # Restart publish task with new resolution
139-
# if self.publish_task:
140-
# self.publish_task.cancel()
141-
# try:
142-
# await self.publish_task
143-
# except asyncio.CancelledError:
144-
# pass
145-
# metadata_cache = LastValueCache[dict]()
146-
# self.publish_task = asyncio.create_task(
147-
# media.run_publish(self.publish_url, self.publish_queue.get, metadata_cache.get, self.emit_monitoring_event)
148-
# )
129+
# TODO: handle prompt changes with differing resolution
149130

150131
logging.info("Received control message with params: %s", data)
151132
yield data

‎runner/app/live/streamer/streamer.py‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from .protocol.protocol import StreamProtocol
1414
from .status import timestamp_to_ms
1515
from trickle import AudioFrame, VideoFrame, OutputFrame, AudioOutput, VideoOutput
16+
from utils import ComfyUtils
1617

1718
fps_log_interval = 10
1819
status_report_interval = 10
@@ -35,15 +36,18 @@ def __init__(
3536
self.tasks_supervisor_task: asyncio.Task | None = None
3637
self.request_id = request_id
3738
self.stream_id = stream_id
38-
self.width = 384
39-
self.height = 704
39+
self.width = ComfyUtils.DEFAULT_WIDTH
40+
self.height = ComfyUtils.DEFAULT_HEIGHT
4041

4142
async def start(self, params: dict):
42-
43-
#TODO: parse from request params
44-
#if params.get('prompt'):
45-
# prompt = params.get('prompt')
46-
# self.width, self.height = ComfyUtils.get_latent_image_dimensions(prompt)
43+
# Parse expected input resolution from workflow prompt
44+
try:
45+
if params.get('prompt'):
46+
prompt = params.get('prompt')
47+
self.width, self.height = ComfyUtils.get_latent_image_dimensions(prompt)
48+
except Exception as e:
49+
logging.error(f"Error parsing resolution from prompt, using default dimensions {ComfyUtils.DEFAULT_WIDTH}x{ComfyUtils.DEFAULT_HEIGHT}: {e}")
50+
self.width, self.height = ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
4751

4852
if self.tasks_supervisor_task:
4953
raise RuntimeError("Streamer already started")

‎runner/app/live/utils/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
from .comfy_utils import ComfyUtils
2+
3+
__all__ = ['ComfyUtils']

‎runner/app/live/utils/comfy_utils.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import logging
2+
3+
class ComfyUtils:
4+
DEFAULT_WIDTH = 512
5+
DEFAULT_HEIGHT = 512
6+
7+
@staticmethod
8+
def get_latent_image_dimensions(workflow: dict | None) -> tuple[int, int]:
9+
"""Get dimensions from the EmptyLatentImage node in the workflow.
10+
11+
Args:
12+
workflow: The workflow JSON dictionary
13+
14+
Returns:
15+
Tuple of (width, height) from the latent image. Returns default dimensions if not found or on error.
16+
"""
17+
18+
if workflow is None:
19+
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT
20+
21+
try:
22+
for node_id, node in workflow.items():
23+
if node.get("class_type") == "EmptyLatentImage":
24+
inputs = node.get("inputs", {})
25+
width = inputs.get("width")
26+
height = inputs.get("height")
27+
if width is not None and height is not None:
28+
return width, height
29+
logging.warning("Incomplete dimensions in latent image node")
30+
break
31+
except Exception as e:
32+
logging.warning(f"Failed to extract dimensions from workflow: {e}")
33+
34+
# Return defaults if dimensions not found or on any error
35+
logging.info(f"Using default dimensions {ComfyUtils.DEFAULT_WIDTH}x{ComfyUtils.DEFAULT_HEIGHT}")
36+
return ComfyUtils.DEFAULT_WIDTH, ComfyUtils.DEFAULT_HEIGHT

0 commit comments

Comments
 (0)