Update myapp.py

myapp.py

@@ -1,5 +1,6 @@
 from flask import Flask, request, jsonify
 from diffusers import DiffusionPipeline
+from transformers import AutoTokenizer
 import torch
 from PIL import Image
 import os
@@ -10,6 +11,9 @@ myapp = Flask(__name__)
 # Load the Diffusion pipeline
 pipe = DiffusionPipeline.from_pretrained("prompthero/openjourney-v4").to("cpu")
 
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained("prompthero/openjourney-v4", clean_up_tokenization_spaces=False)
+
 @myapp.route('/')
 def index():
     return '''
@@ -53,12 +57,15 @@ def generate_image():
     data = request.json
     prompt = data.get('prompt', 'Astronaut in a jungle, cold color palette, muted colors, detailed, 8k')
 
+    # Tokenize the prompt
+    tokens = tokenizer(prompt, return_tensors="pt")
+
     # Generate the image
     image = pipe(prompt).images[0]
 
     # Convert to PIL Image and save
     pil_image = Image.fromarray(image.numpy())
-    output_path = f"{prompt.replace(' ', '_')}.png"  # Create a file name based on the prompt
+    output_path = f"{prompt.replace(' ', '_')}.png"
     pil_image.save(output_path)
 
     # Return the path to the generated image