app.py

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_prompt(starting_text):
    for count in range(4):
        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=4)
        response_list = []
        for x in response:
            resp = x['generated_text'].strip()
            if resp != starting_text and len(resp) > (len(starting_text) + 4) and not resp.endswith((":", "-", "—")):
                response_list.append(resp + '\n')

        # Clean the generated text
        response_end = "\n".join(response_list)
        response_end = re.sub(r'[^ ]+\.[^ ]+', '', response_end)  # Removes strings like 'abc.xyz'
        response_end = response_end.replace("<", "").replace(">", "")

        if response_end:
            return response_end
        if count == 4:
            return response_end

# Define the API endpoint
@app.route('/', methods=['GET'])
def generate_api():
    starting_text = request.args.get('text', default="", type=str)
    result = generate_prompt(starting_text)
    return jsonify({"generated_text": result})

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