Create app.py

app.py

@@ -0,0 +1,235 @@
+# This script demonstrates the use of the websockets API and the SaveImageWebsocket node to retrieve images directly
+# without saving them to disk. It includes a Gradio interface for user interaction.
+
+# JSON Comment Section
+"""
+{
+    "inputs": [
+        {
+            "type": "text",
+            "label": "Text Prompt",
+            "name": "text_prompt",
+            "default": "masterpiece best quality man"
+        },
+        {
+            "type": "number",
+            "label": "Seed",
+            "name": "seed",
+            "default": 5
+        },
+        {
+            "type": "radio",
+            "label": "Server",
+            "name": "server",
+            "choices": ["AWS Server", "Home Server"],
+            "default": "Home Server"
+        }
+    ],
+    "outputs": [
+        {
+            "type": "image",
+            "label": "Generated Image",
+            "name": "generated_image"
+        }
+    ]
+}
+"""
+
+import gradio as gr
+import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+import uuid
+import json
+import urllib.request
+import urllib.parse
+from PIL import Image
+import io
+
+# Generate a unique client ID for the session
+client_id = str(uuid.uuid4())
+
+# Function to queue a prompt to the server
+def queue_prompt(prompt, server_address):
+    p = {"prompt": prompt, "client_id": client_id}
+    data = json.dumps(p).encode('utf-8')
+    req = urllib.request.Request("http://{}/prompt".format(server_address), data=data)
+    return json.loads(urllib.request.urlopen(req).read())
+
+# Function to retrieve an image from the server
+def get_image(filename, subfolder, folder_type, server_address):
+    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    url_values = urllib.parse.urlencode(data)
+    with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response:
+        return response.read()
+
+# Function to retrieve the history of a prompt from the server
+def get_history(prompt_id, server_address):
+    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
+        return json.loads(response.read())
+
+# Function to get images from the websocket connection
+def get_images(ws, prompt, server_address):
+    prompt_id = queue_prompt(prompt, server_address)['prompt_id']
+    output_images = {}
+    current_node = ""
+    while True:
+        out = ws.recv()
+        if isinstance(out, str):
+            message = json.loads(out)
+            if message['type'] == 'executing':
+                data = message['data']
+                if data['prompt_id'] == prompt_id:
+                    if data['node'] is None:
+                        break # Execution is done
+                    else:
+                        current_node = data['node']
+        else:
+            if current_node == 'save_image_websocket_node':
+                images_output = output_images.get(current_node, [])
+                images_output.append(out[8:])
+                output_images[current_node] = images_output
+
+    return output_images
+
+# Function to generate an image based on the provided text prompt, seed, and server
+def generate_image(text_prompt, seed, server):
+    prompt_text = """
+    {
+        "3": {
+            "class_type": "KSampler",
+            "inputs": {
+                "cfg": 8,
+                "denoise": 1,
+                "latent_image": [
+                    "5",
+                    0
+                ],
+                "model": [
+                    "4",
+                    0
+                ],
+                "negative": [
+                    "7",
+                    0
+                ],
+                "positive": [
+                    "6",
+                    0
+                ],
+                "sampler_name": "euler",
+                "scheduler": "normal",
+                "seed": 8566257,
+                "steps": 20
+            }
+        },
+        "4": {
+            "class_type": "CheckpointLoaderSimple",
+            "inputs": {
+                "ckpt_name": "v1-5-pruned-emaonly.ckpt"
+            }
+        },
+        "5": {
+            "class_type": "EmptyLatentImage",
+            "inputs": {
+                "batch_size": 1,
+                "height": 512,
+                "width": 768
+            }
+        },
+        "6": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {
+                "clip": [
+                    "4",
+                    1
+                ],
+                "text": "masterpiece best quality girl"
+            }
+        },
+        "7": {
+            "class_type": "CLIPTextEncode",
+            "inputs": {
+                "clip": [
+                    "4",
+                    1
+                ],
+                "text": "bad hands"
+            }
+        },
+        "8": {
+            "class_type": "VAEDecode",
+            "inputs": {
+                "samples": [
+                    "3",
+                    0
+                ],
+                "vae": [
+                    "4",
+                    2
+                ]
+            }
+        },
+        "save_image_websocket_node": {
+            "class_type": "SaveImageWebsocket",
+            "inputs": {
+                "images": [
+                    "8",
+                    0
+                ]
+            }
+        }
+    }
+    """
+
+    prompt = json.loads(prompt_text)
+    # Set the text prompt for our positive CLIPTextEncode
+    prompt["6"]["inputs"]["text"] = text_prompt
+
+    # Set the seed for our KSampler node
+    prompt["3"]["inputs"]["seed"] = seed
+
+    # Determine the server address based on the selected server
+    server_address = "3.14.144.23:8188" if server == "AWS Server" else "192.168.50.136:8188"
+
+    # Establish a websocket connection
+    ws = websocket.WebSocket()
+    ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
+    images = get_images(ws, prompt, server_address)
+
+    # Get the first image from the output
+    image = None
+
+    for node_id in images:
+        for image_data in images[node_id]:
+            image = Image.open(io.BytesIO(image_data))
+            break
+        if image:
+            break
+
+    return image
+
+# Function to handle cancel request
+def cancel_request():
+    return "Request Cancelled"
+
+# Gradio Interface using Blocks
+with gr.Blocks() as demo:
+    gr.Markdown("# Image Generation with Websockets API")
+    gr.Markdown("Generate images using a Websockets API and SaveImageWebsocket node.")
+    
+    with gr.Row():
+        with gr.Column():
+            text_prompt = gr.Textbox(label="Text Prompt", value="masterpiece best quality man")
+            seed = gr.Number(label="Seed", value=5)
+            server = gr.Radio(label="Server", choices=["AWS Server", "Home Server"], value="Home Server")
+            generate_button = gr.Button("Generate Image")
+            cancel_button = gr.Button("Cancel Request")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+    
+    # Set up button click events
+    generate_button.click(fn=generate_image, inputs=[text_prompt, seed, server], outputs=output_image)
+    cancel_button.click(fn=cancel_request, inputs=[], outputs=[])
+
+# Launch the Gradio interface
+demo.launch()