from transformers import pipeline, set_seed
from flask import Flask, request, jsonify
import random, re

app = Flask(__name__)

# Initialize the GPT-2 pipeline
gpt2_pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion', tokenizer='gpt2')
with open("ideas.txt", "r") as f:
    lines = f.readlines()

def generate_prompts(starting_text, num_prompts=1):
    response_list = []
    
    for _ in range(num_prompts):
        for count in range(4):  # Attempt up to 4 times to generate valid response
            seed = random.randint(100, 1000000)
            set_seed(seed)

            # Choose a random line from the file if the input text is empty
            if starting_text == "":
                starting_text = lines[random.randrange(0, len(lines))].strip().lower().capitalize()
                starting_text = re.sub(r"[,:\-–.!;?_]", '', starting_text)

            # Generate text
            response = gpt2_pipe(starting_text, max_length=random.randint(60, 90), num_return_sequences=1)
            generated_text = response[0]['generated_text'].strip()

            # Clean and check the generated response
            if generated_text != starting_text and len(generated_text) > (len(starting_text) + 4):
                cleaned_text = re.sub(r'[^ ]+\.[^ ]+', '', generated_text)  # Remove strings like 'abc.xyz'
                cleaned_text = cleaned_text.replace("<", "").replace(">", "")
                response_list.append(cleaned_text)
                break  # Stop trying further once a valid prompt is added

    return response_list[:num_prompts]

# Define the API endpoint
@app.route('/', methods=['GET'])
def generate_api():
    starting_text = request.args.get('text', default="", type=str)
    num_prompts = request.args.get('n', default=1, type=int)  # Get the number of prompts to return, default is 1
    
    # Generate the prompts
    results = generate_prompts(starting_text, num_prompts=num_prompts)
    return jsonify(results)

if __name__ == '__main__':
    # Run the Flask app on port 7860
    app.run(host='0.0.0.0', port=7860)