Update app.py

app.py

@@ -41,22 +41,23 @@ def wait_for_image_with_prefix(folder, prefix):
 
     while True:
         files = os.listdir(folder)
-        image_files = [f for f in files if f.lower().startswith(prefix.lower()) and 
+        image_files = [f for f in files if f.lower().startswith(prefix.lower()) and
                        f.lower().endswith(('.png', '.jpg', '.jpeg'))]
-        
+
         if image_files:
             # Sort by modification time to get the latest file
             image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder, x)), reverse=True)
             latest_image = os.path.join(folder, image_files[0])
-            
+
             if is_file_ready(latest_image):
                 # Wait a bit more to ensure the file is completely written
                 time.sleep(1)
                 return latest_image
-        
+
         # If no matching file found, wait before checking again
         time.sleep(1)
 
+
 def delete_image_file(file_path):
     try:
         if os.path.exists(file_path):
@@ -67,6 +68,7 @@ def delete_image_file(file_path):
     except Exception as e:
         logger.debug(f"error {file_path}: {str(e)}")
 
+
 def start_queue(prompt_workflow):
     p = {"prompt": prompt_workflow}
     data = json.dumps(p).encode('utf-8')
@@ -80,16 +82,19 @@ def check_server_ready():
     except requests.RequestException:
         return False
 
-queue_reqs=set()
+
+queue_reqs = set()
+
 
 @spaces.GPU(duration=240)
 def generate_image(prompt, image):
+    prefix_filename = str(random.randint(0, 999999))
+    prompt = prompt.replace('ComfyUI', prefix_filename)
     prompt = json.loads(prompt)
-    
+
     image = Image.fromarray(image)
-    image.save(INPUT_DIR+'/input.png', format='PNG')
-    
-    prefix_filename = str(random.randint(0, 999999))
+    image.save(INPUT_DIR + '/input.png', format='PNG')
+
     queue_reqs.add(prefix_filename)
     process = None
     try:
@@ -115,10 +120,7 @@ def generate_image(prompt, image):
             if latest_image:
                 logger.debug(f"file is: {latest_image}")
                 try:
-                    with open(latest_image, 'rb') as f:
-                        photo = f.read()
-                    logger.debug(f"file bytes size: {len(photo)}")
-                    return io.BytesIO(photo)
+                    return Image.open(latest_image)
                 finally:
                     delete_image_file(latest_image)
             time.sleep(1)
@@ -139,16 +141,14 @@ def generate_image(prompt, image):
                 process.kill()
 
 
-
-
 if __name__ == "__main__":
     demo = gr.Interface(fn=generate_image, inputs=[
         "text",
         gr.Image(image_mode='RGBA', type="numpy")
     ],
-    outputs=[
-        gr.Image(type="numpy", image_mode='RGBA')
-    ])
+                        outputs=[
+                            gr.Image(type="numpy", image_mode='RGBA')
+                        ])
     demo.launch(debug=True)
     logger.debug('demo.launch()')