ltdrdata / ComfyUI-Workflow-Component

This is a side project to experiment with using workflows as components.
GNU General Public License v3.0
214 stars 9 forks source link

Enhance 'prompt_hook' for better key handling #10

Closed lcolok closed 1 year ago

lcolok commented 1 year ago

I hope this message finds you well.

Recently, while trying out the API request example code found at this location, I encountered a KeyError: 'extra_data'. Upon investigation, I noticed that the current implementation of the prompt_hook function directly accesses the dictionary keys 'extra_data', 'extra_pnginfo', 'workflow', and 'nodes', which could lead to a KeyError if these keys do not exist.

To address this, I am submitting this Pull Request, offering an improvement to this section of the code. I suggest utilizing the .get() method for dictionary access to retrieve these keys safely. The .get() method will return a default value if the keys do not exist, thereby avoiding the KeyError.

Here is the modified section of the code:

extra_data = json_data.get('extra_data', {})
extra_pnginfo = extra_data.get('extra_pnginfo', {})
workflow = extra_pnginfo.get('workflow', {})
nodes = workflow.get('nodes', [])

I believe this small adjustment can significantly enhance the stability and reliability of the ComfyUI project when handling unpredictable requests.

Once again, thank you for your tremendous contributions to this project. I look forward to your feedback on this proposed enhancement.