add image to image workflow

Author: phDoll

.gitignore

@@ -1,4 +1,5 @@
 output*
 workflow*
+input*
 __pycache__/
 .DS_Store*

api/websocket_api.py

@@ -9,14 +9,30 @@
 from PIL import Image
 import io
 from utils.helpers.find_node import find_node
-
+from requests_toolbelt import MultipartEncoder
 
 server_address='127.0.0.1:8188'
 client_id=str(uuid.uuid4())
 
 ws = websocket.WebSocket()
 ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
 
+def setup_image(input_path, name, type="input", overwrite=False):
+    with open(input_path, 'rb') as file:
+        multipart_data = MultipartEncoder(
+            fields={
+                'image': (name, file, 'image/jpeg'),  # Adjust the content-type accordingly
+                'type': type,
+                'overwrite': str(overwrite).lower()
+            }
+        )
+
+        data = multipart_data
+        headers = {'Content-Type': multipart_data.content_type}
+        request = urllib.request.Request("http://{}/upload/image".format(server_address), data=data, headers=headers)
+        with urllib.request.urlopen(request) as response:
+            return response.read()
+
 def generate_image_by_prompt(prompt, output_path, save_previews=False):
     prompt_id = queue_prompt(prompt)['prompt_id']
     track_progress(prompt, ws, prompt_id)

main.py

@@ -1,4 +1,5 @@
 from utils.actions.prompt_to_image import prompt_to_image
+from utils.actions.prompt_image_to_image import prompt_image_to_image
 from utils.actions.load_workflow import load_workflow
 from api.websocket_api import clear_comfy_cache
 from api.websocket_api import get_image
@@ -8,9 +9,11 @@
 def main():
     try:
       print("Welcome to the program!")
-      workflow = load_workflow('./workflows/image_refeed_workflow.json')
+      workflow = load_workflow('./workflows/image_to_image.json')
 
-      prompt_to_image(workflow, 'beautiful woman sitting on a desk in a nice restaurant, candlelight dinner atmosphere, wearing a red dress', save_previews=True)
+      # prompt_to_image(workflow, 'beautiful woman sitting on a desk in a nice restaurant, candlelight dinner atmosphere, wearing a red dress', save_previews=True)
+      input_path = './input/ComfyUI_00103_.png'
+      prompt_image_to_image(workflow, input_path, 'beautiful [white woman], (dark lighting), curly blond hair', save_previews=True)
     except Exception as e:
       print(f"An error occurred: {e}")
       exit_program()

requirements.txt

@@ -1,3 +1,4 @@
 Pillow==10.0.0
 Pillow==10.2.0
+requests_toolbelt==1.0.0
 websocket_client==1.7.0

utils/actions/prompt_image_to_image.py

@@ -0,0 +1,33 @@
+from api.websocket_api import generate_image_by_prompt
+from api.websocket_api import setup_image
+from utils.helpers.find_node import find_node
+from utils.helpers.randomize_seed import generate_random_15_digit_number
+from utils.helpers.replace_key import replace_key
+from utils.helpers.find_parent import find_parent_of_key
+import json
+def prompt_image_to_image(workflow, input_path, positve_prompt, negative_prompt='', save_previews=False):
+  prompt = json.loads(workflow)
+  replace_key(prompt, 'seed', generate_random_15_digit_number())
+  id_to_class_type = {id: details['class_type'] for id, details in prompt.items()}
+  k_sampler = [key for key, value in id_to_class_type.items() if value == 'KSampler'][0]
+
+  postive_input_id = prompt.get(k_sampler)['inputs']['positive'][0]
+  prompt.get(postive_input_id)['inputs']['text_g'] = positve_prompt
+  prompt.get(postive_input_id)['inputs']['text_l'] = positve_prompt
+
+  if negative_prompt != '':
+    negative_input_id = prompt.get(k_sampler)['inputs']['negative'][0]
+    id_to_class_type.get(negative_input_id)['inputs']['text_g'] = negative_prompt
+    id_to_class_type.get(negative_input_id)['inputs']['text_l'] = negative_prompt
+
+
+  image_loader = [key for key, value in id_to_class_type.items() if value == 'LoadImage'][0]
+  filename = input_path.split('/')[-1]
+  prompt.get(image_loader)['inputs']['image'] = filename
+
+  setup_image(input_path, filename)
+  generate_image_by_prompt(prompt, './output/', save_previews)
+
+
+
+  

utils/helpers/find_parent.py

@@ -0,0 +1,10 @@
+def find_parent_of_key(json_obj, key, parent=None):
+    if key in json_obj:
+        return parent
+    for k, v in json_obj.items():
+        if isinstance(v, dict):
+            # Pass the current object as the parent for the next level of recursion
+            parent_obj = find_parent_of_key(v, key, parent=json_obj)
+            if parent_obj is not None:
+                return parent_obj
+    return None