Update app.py

app.py

@@ -1,15 +1,15 @@
 import os
-from flask import Flask, request, jsonify, send_file, render_template
-import websocket
-import uuid
+import io
 import json
+import base64
+import random
 import urllib.request
 import urllib.parse
-import random
-from PIL import Image
-import io
-import base64
+import websocket
+import uuid
 from dotenv import load_dotenv
+from flask import Flask, request, jsonify, render_template
+from PIL import Image
 
 # Load environment variables from the .env file
 load_dotenv()
@@ -17,18 +17,11 @@ load_dotenv()
 # Initialize Flask app
 app = Flask(__name__)
 
-# Set server and websocket addresses from environment variables
-# server_address = "https://gosign-de-comfyui.hf.space"
-# ws_address = "wss://gosign-de-comfyui.hf.space/ws"
-
 # Set server and websocket addresses from environment variables
 server_address = os.getenv("SERVER_ADDRESS")
 ws_address = os.getenv("WS_ADDRESS")
 
-# Example print statements to verify the addresses are loaded
-print(f"Server Address: {server_address}")
-print(f"WebSocket Address: {ws_address}")
-
+# Generate a unique client ID
 client_id = str(uuid.uuid4())
 
 def make_request(url, data=None, headers=None):
@@ -40,7 +33,7 @@ def queue_prompt(prompt, token):
     payload = {"prompt": prompt, "client_id": client_id}
     data = json.dumps(payload).encode('utf-8')
     headers = {
-        'Authorization': f'Bearer {token}',  # Bearer token
+        'Authorization': f'Bearer {token}',
         'Content-Type': 'application/json'
     }
     return make_request(f"{server_address}/prompt", data=data, headers=headers)
@@ -48,19 +41,13 @@ def queue_prompt(prompt, token):
 def get_image(filename, subfolder, image_type, token):
     url_values = {'filename': filename, 'subfolder': subfolder, 'type': image_type}
     url = f"{server_address}/view?{urllib.parse.urlencode(url_values)}"
-
     req = urllib.request.Request(url)
     req.add_header("Authorization", f"Bearer {token}")
-
-    # Debugging output
-    print(f"Request URL: {url}")
-    print(f"Request Headers: {req.headers}")
-
     try:
         return urllib.request.urlopen(req).read()
     except urllib.error.HTTPError as e:
         print(f"HTTP Error: {e.code} - {e.reason}")
-        print(e.read())  # This will give you more context about the error
+        print(e.read())
         raise
 
 def get_history(prompt_id, token):
@@ -94,6 +81,7 @@ def get_images(ws, prompt, token):
         output_images[node_id] = images_output
 
     return output_images
+
 # Default route for home welcome
 @app.route('/')
 def home():
@@ -107,14 +95,10 @@ def generate_image():
     # Extract the token from the request headers
     token = request.headers.get('Authorization')
 
-    # Ensure the token is present and in the correct format
     if token is None:
         return jsonify({'error': 'No token provided'}), 400
-    # If the token does not have a 'Bearer' prefix, use the entire token
     if token.startswith("Bearer "):
         token = token.split(" ")[1]  # Extract the actual token if 'Bearer ' exists
-    else:
-        token = token  # Use the entire token if there's no 'Bearer ' prefix
 
     # Base64 decode the encoded token
     token = base64.b64decode(token).decode("utf-8")
@@ -124,8 +108,6 @@ def generate_image():
 
     text_prompt = data['text_prompt']
     
-    # file_path = '/Users/muhammadwaqas/gosign/comfyui-api/workflow_api.json'
-
     # Get the path to the current file's directory
     current_dir = os.path.dirname(os.path.abspath(__file__))
     file_path = os.path.join(current_dir, 'workflow_api.json')
@@ -145,17 +127,17 @@ def generate_image():
     images = get_images(ws, prompt, token)
     ws.close()
 
-    output_image_paths = []
+    output_images_base64 = []
 
     for node_id in images:
         for image_data in images[node_id]:
             image = Image.open(io.BytesIO(image_data))
-            image_path = f'image-output-{seednum}-{node_id}.png'
-            image.show()
-            image.save(image_path)
-            output_image_paths.append(image_path)
+            buffered = io.BytesIO()
+            image.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+            output_images_base64.append(img_str)
 
-    return jsonify({'images': output_image_paths})
+    return jsonify({'images': output_images_base64})
 
 @app.route('/get_image/<filename>', methods=['GET'])
 def get_image_file(filename):