Update app.py

app.py

@@ -1,42 +1,7 @@
-# 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 subprocess
+import threading
 import gradio as gr
-import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+import websocket
 import uuid
 import json
 import urllib.request
@@ -44,29 +9,44 @@ import urllib.parse
 from PIL import Image
 import io
 
-# Generate a unique client ID for the session
+# 🚀 Chapter 1: Install Necessary Packages 🚀
+def install_packages():
+    packages = [
+        "gradio",
+        "websocket-client",
+        "pillow"
+    ]
+    for package in packages:
+        subprocess.check_call(["pip", "install", package])
+
+# Use threading to run the installation in the background
+install_thread = threading.Thread(target=install_packages)
+install_thread.start()
+install_thread.join()
+
+# 🌟 Chapter 2: Generate Client ID 🌟
 client_id = str(uuid.uuid4())
 
-# Function to queue a prompt to the server
+# 🌟 Chapter 3: Queue Prompt Function 🌟
 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)
+    req = urllib.request.Request(f"http://{server_address}/prompt", data=data)
     return json.loads(urllib.request.urlopen(req).read())
 
-# Function to retrieve an image from the server
+# 🌟 Chapter 4: Get Image Function 🌟
 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:
+    with urllib.request.urlopen(f"http://{server_address}/view?{url_values}") as response:
         return response.read()
 
-# Function to retrieve the history of a prompt from the server
+# 🌟 Chapter 5: Get History Function 🌟
 def get_history(prompt_id, server_address):
-    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
+    with urllib.request.urlopen(f"http://{server_address}/history/{prompt_id}") as response:
         return json.loads(response.read())
 
-# Function to get images from the websocket connection
+# 🌟 Chapter 6: Get Images Function 🌟
 def get_images(ws, prompt, server_address):
     prompt_id = queue_prompt(prompt, server_address)['prompt_id']
     output_images = {}
@@ -79,7 +59,7 @@ def get_images(ws, prompt, server_address):
                 data = message['data']
                 if data['prompt_id'] == prompt_id:
                     if data['node'] is None:
-                        break # Execution is done
+                        break
                     else:
                         current_node = data['node']
         else:
@@ -90,7 +70,7 @@ def get_images(ws, prompt, server_address):
 
     return output_images
 
-# Function to generate an image based on the provided text prompt, seed, and server
+# 🌟 Chapter 7: Generate Image Function 🌟
 def generate_image(text_prompt, seed, server):
     prompt_text = """
     {
@@ -181,21 +161,12 @@ def generate_image(text_prompt, seed, server):
     """
 
     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))
+    ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
     images = get_images(ws, prompt, server_address)
-
-    # Get the first image from the output
     image = None
 
     for node_id in images:
@@ -207,11 +178,11 @@ def generate_image(text_prompt, seed, server):
 
     return image
 
-# Function to handle cancel request
+# 🌟 Chapter 8: Cancel Request Function 🌟
 def cancel_request():
     return "Request Cancelled"
 
-# Gradio Interface using Blocks
+# 🌟 Chapter 9: Gradio Interface 🌟
 with gr.Blocks() as demo:
     gr.Markdown("# Image Generation with Websockets API")
     gr.Markdown("Generate images using a Websockets API and SaveImageWebsocket node.")
@@ -227,9 +198,7 @@ with gr.Blocks() as demo:
         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()