@@ -51,6 +51,49 @@ def validate_prompt(cls, v) -> dict:
5151 raise ValueError ("Prompt must be either a JSON object or such JSON object serialized as a string" )
5252
5353
54+ class ComfyUtils :
55+ @staticmethod
56+ def get_latent_image_dimensions (workflow : dict ) -> tuple [int , int ]:
57+ """Get dimensions from the EmptyLatentImage node in the workflow.
58+
59+ Args:
60+ workflow: The workflow JSON dictionary
61+
62+ Returns:
63+ Tuple of (width, height) from the latent image, or (None, None) if not found
64+ """
65+ for node_id , node in workflow .items ():
66+ if node .get ("class_type" ) == "EmptyLatentImage" :
67+ try :
68+ inputs = node .get ("inputs" , {})
69+ return inputs .get ("width" ), inputs .get ("height" )
70+ except Exception as e :
71+ logging .warning (f"Failed to extract dimensions from latent image: { e } )
72+ return None , None
73+ return None , None
74+
75+ @staticmethod
76+ def update_latent_image_dimensions (workflow : dict , width : int , height : int ) -> dict | None :
77+ """Update the EmptyLatentImage node dimensions in the workflow.
78+
79+ Args:
80+ workflow: The workflow JSON dictionary
81+ width: Width to set
82+ height: Height to set
83+ """
84+ for node_id , node in workflow .items ():
85+ if node .get ("class_type" ) == "EmptyLatentImage" :
86+ try :
87+ if "inputs" not in node :
88+ node ["inputs" ] = {}
89+ node ["inputs" ]["width" ] = width
90+ node ["inputs" ]["height" ] = height
91+ logging .info (f"Updated latent image dimensions to { width } { height } )
92+ except Exception as e :
93+ logging .warning (f"Failed to update latent image dimensions: { e } )
94+ break
95+
96+
5497class ComfyUI (Pipeline ):
5598 def __init__ (self ):
5699 comfy_ui_workspace = os .getenv (COMFY_UI_WORKSPACE_ENV )
@@ -61,14 +104,30 @@ def __init__(self):
61104 async def initialize (self , ** params ):
62105 new_params = ComfyUIParams (** params )
63106 logging .info (f"Initializing ComfyUI Pipeline with prompt: { new_params .prompt } )
64- # TODO: currently its a single prompt, but need to support multiple prompts
107+
108+ # Get dimensions from workflow if it's a dict
109+
110+ if width is None or height is None :
111+ if isinstance (new_params .prompt , dict ):
112+ # If dimensions not provided in params, get them from latent image
113+ latent_width , latent_height = ComfyUtils .get_latent_image_dimensions (new_params .prompt )
114+ new_params .width = width or latent_width or new_params .width
115+ new_params .height = height or latent_height or new_params .height
116+ else :
117+ # If dimensions provided in params, update the latent image
118+ ComfyUtils .update_latent_image_dimensions (new_params .prompt , width , height )
119+
120+ # TODO clean up extra vars
121+ width = width or new_params .width
122+ height = height or new_params .height
123+
65124 await self .client .set_prompts ([new_params .prompt ])
66125 self .params = new_params
67126
68- # Warm up the pipeline
69- logging .info (f"Warming up pipeline with dimensions: { new_params . width } { new_params . height } )
127+ # Warm up the pipeline with the final dimensions
128+ logging .info (f"Warming up pipeline with dimensions: { width } { height } )
70129 dummy_frame = VideoFrame (None , 0 , 0 )
71- dummy_frame .side_data .input = torch .randn (1 , new_params . height , new_params . width , 3 )
130+ dummy_frame .side_data .input = torch .randn (1 , height , width , 3 )
72131
73132 for _ in range (WARMUP_RUNS ):
74133 self .client .put_video_input (dummy_frame )
0 commit comments