浏览代码

add image name to output

Sergiu 9 月之前
父节点
当前提交
0316f014d4
共有 1 个文件被更改,包括 14 次插入3 次删除
  1. 14 3
      04_stable_diffusion/sd_comfy_api_v2.py

+ 14 - 3
04_stable_diffusion/sd_comfy_api_v2.py

@@ -131,9 +131,9 @@ def get_images(ws, prompt):
             if message['type'] == 'executing':
                 data = message['data']
                 if data['node'] is None and data['prompt_id'] == prompt_id:
-                    break #Execution is done
+                    break  # Execution is done
         else:
-            continue #previews are binary data
+            continue  # previews are binary data
 
     history = get_history(prompt_id)[prompt_id]
     for node_id in history['outputs']:
@@ -142,7 +142,11 @@ def get_images(ws, prompt):
         if 'images' in node_output:
             for image in node_output['images']:
                 image_data = get_image(image['filename'], image['subfolder'], image['type'])
-                images_output.append(image_data)
+                images_output.append({
+                    'filename': image['filename'],
+                    'data': image_data,
+                    'type': image['type']
+                })
         output_images[node_id] = images_output
 
     return output_images
@@ -172,6 +176,13 @@ def main():
 
     prompt_id = queue_prompt(prompt)['prompt_id']
 
+    for node_id in images:
+        for image_info in images[node_id]:
+            if image_info['type'] == 'output':
+                print("Upload image")
+                
+            
+