Updates

2026-01-28 16:54:06 +01:00
parent 213c2836f9
commit 493eb0ed90
73 changed files with 0 additions and 0 deletions
--- a/DreamCanvas/main.py
+++ b/DreamCanvas/main.py
@@ -0,0 +1,188 @@
+import websocket  # websocket-client
+import uuid
+import json
+import urllib.request
+import urllib.parse
+import random
+from PIL import Image
+import io
+from termcolor import colored
+from dotenv import load_dotenv
+import os
+
+# Step 1: Initialize the connection settings and load environment variables
+print(colored("Step 1: Initialize the connection settings and load environment variables.", "cyan"))
+print(colored("Loading configuration from the .env file.", "yellow"))
+load_dotenv()
+
+# Get server address from environment variable, default to "localhost:8188" if not set
+server_address = os.getenv('COMFYUI_SERVER_ADDRESS', 'localhost:8188')
+client_id = str(uuid.uuid4())
+
+# Display the server address and client ID for transparency
+print(colored(f"Server Address: {server_address}", "magenta"))
+print(colored(f"Generated Client ID: {client_id}", "magenta"))
+input(colored("Press Enter to continue...", "green"))
+
+# Queue prompt function
+def queue_prompt(prompt):
+    p = {"prompt": prompt, "client_id": client_id}
+    data = json.dumps(p, indent=4).encode('utf-8')  # Prettify JSON for print
+    req = urllib.request.Request(f"http://{server_address}/prompt", data=data)
+    
+    # Step 5: Queue the prompt and prepare to send it to the ComfyUI server
+    print(colored(f"Step 5: Queuing the prompt for client ID {client_id}.", "cyan"))
+    input(colored("Press Enter to view the JSON that will be sent...", "green"))
+    
+    # Pretty-printed JSON for the prompt
+    print(colored("Here's the JSON that will be sent:", "yellow"))
+    print(colored(json.dumps(p, indent=4), "blue"))  # Pretty-printed JSON
+    
+    input(colored("Press Enter to continue and send the prompt...", "green"))
+    
+    return json.loads(urllib.request.urlopen(req).read())
+
+# Get image function
+def get_image(filename, subfolder, folder_type):
+    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    url_values = urllib.parse.urlencode(data)
+    
+    print(colored(f"Fetching image from the server: {server_address}/view", "cyan"))
+    print(colored(f"Filename: {filename}, Subfolder: {subfolder}, Type: {folder_type}", "yellow"))
+    with urllib.request.urlopen(f"http://{server_address}/view?{url_values}") as response:
+        return response.read()
+
+# Get history for a prompt ID
+def get_history(prompt_id):
+    print(colored(f"Fetching history for prompt ID: {prompt_id}.", "cyan"))
+    with urllib.request.urlopen(f"http://{server_address}/history/{prompt_id}") as response:
+        return json.loads(response.read())
+
+# Get images from the workflow
+def get_images(ws, prompt):
+    prompt_id = queue_prompt(prompt)['prompt_id']
+    output_images = {}
+
+    last_reported_percentage = 0
+    
+    print(colored("Step 6: Start listening for progress updates via the WebSocket connection.", "cyan"))
+    input(colored("Press Enter to continue...", "green"))
+
+    while True:
+        out = ws.recv()
+        if isinstance(out, str):
+            message = json.loads(out)
+            if message['type'] == 'progress':
+                data = message['data']
+                current_progress = data['value']
+                max_progress = data['max']
+                percentage = int((current_progress / max_progress) * 100)
+
+                # Only update progress every 10%
+                if percentage >= last_reported_percentage + 10:
+                    print(colored(f"Progress: {percentage}% in node {data['node']}", "yellow"))
+                    last_reported_percentage = percentage
+
+            elif message['type'] == 'executing':
+                data = message['data']
+                if data['node'] is None and data['prompt_id'] == prompt_id:
+                    print(colored("Execution complete.", "green"))
+                    break  # Execution is done
+        else:
+            continue  # Previews are binary data
+
+    # Fetch history and images after completion
+    print(colored("Step 7: Fetch the history and download the images after execution completes.", "cyan"))
+    input(colored("Press Enter to continue...", "green"))
+
+    history = get_history(prompt_id)[prompt_id]
+    for o in history['outputs']:
+        for node_id in history['outputs']:
+            node_output = history['outputs'][node_id]
+            if 'images' in node_output:
+                images_output = []
+                for image in node_output['images']:
+                    print(colored(f"Downloading image: {image['filename']} from the server.", "yellow"))
+                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
+                    images_output.append(image_data)
+                output_images[node_id] = images_output
+
+    return output_images
+
+# Generate images function with customizable input
+def generate_images(positive_prompt, negative_prompt, steps=25, resolution=(512, 512)):
+    # Step 3: Establish WebSocket connection
+    ws = websocket.WebSocket()
+    ws_url = f"ws://{server_address}/ws?clientId={client_id}"
+    print(colored(f"Step 3: Establishing WebSocket connection to {ws_url}", "cyan"))
+    input(colored("Press Enter to continue...", "green"))
+    ws.connect(ws_url)
+    
+    # Step 4: Load workflow from file and print it
+    print(colored("Step 4: Loading the image generation workflow from 'workflow.json'.", "cyan"))
+    with open("workflow.json", "r", encoding="utf-8") as f:
+        workflow_data = f.read()
+
+    workflow = json.loads(workflow_data)
+    
+    input(colored("Press Enter to view the loaded workflow before customization...", "green"))
+    
+    # Print the loaded workflow before customization
+    print(colored("Here's the workflow as it was loaded before customization:", "yellow"))
+    print(colored(json.dumps(workflow, indent=4), "blue"))  # Pretty-print the workflow
+    
+    input(colored("Press Enter to continue to customization...", "green"))
+
+    # Customize workflow based on inputs
+    print(colored("Step 5: Customizing the workflow with the provided inputs.", "cyan"))
+    print(colored(f"Setting positive prompt: {positive_prompt}", "yellow"))
+    print(colored(f"Setting negative prompt: {negative_prompt}", "yellow"))
+    workflow["6"]["inputs"]["text"] = positive_prompt
+    workflow["7"]["inputs"]["text"] = negative_prompt
+
+    print(colored(f"Setting steps for generation: {steps}", "yellow"))
+    workflow["3"]["inputs"]["steps"] = steps
+
+    print(colored(f"Setting resolution to {resolution[0]}x{resolution[1]}", "yellow"))
+    workflow["5"]["inputs"]["width"] = resolution[0]
+    workflow["5"]["inputs"]["height"] = resolution[1]
+
+    # Set a random seed for the KSampler node
+    seed = random.randint(1, 1000000000)
+    print(colored(f"Setting random seed for generation: {seed}", "yellow"))
+    workflow["3"]["inputs"]["seed"] = seed
+    
+    input(colored("Press Enter to continue...", "green"))
+
+    # Fetch generated images
+    images = get_images(ws, workflow)
+
+    # Step 8: Close WebSocket connection after fetching the images
+    print(colored(f"Step 8: Closing WebSocket connection to {ws_url}", "cyan"))
+    ws.close()
+    input(colored("Press Enter to continue...", "green"))
+
+    return images, seed
+
+# Example of calling the method and saving the images
+if __name__ == "__main__":
+    # Step 2: User input for prompts
+    positive_prompt = input(colored("Enter the positive prompt: ", "cyan"))
+    negative_prompt = input(colored("Enter the negative prompt: ", "cyan"))
+
+    print(colored("Step 2: User inputs the positive and negative prompts for image generation.", "cyan"))
+    input(colored("Press Enter to continue...", "green"))
+
+    # Call the generate_images function
+    images, seed = generate_images(positive_prompt, negative_prompt)
+
+    # Step 9: Save the images
+    print(colored("Step 9: Saving the generated images locally.", "cyan"))
+    input(colored("Press Enter to continue...", "green"))
+    
+    for node_id in images:
+        for image_data in images[node_id]:
+            image = Image.open(io.BytesIO(image_data))
+            filename = f"outputs/{node_id}-{seed}.png"
+            image.save(filename)
+            print(colored(f"Image saved as {filename}", "blue"))