from flask import Flask, render_template, request


import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "cagliostrolab/animagine-xl-3.1", 
    torch_dtype=torch.float16, 
    use_safetensors=True, 
)
pipe.to('cpu')




app = Flask(__name__)



@app.route('/')
def home():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate():
    if request.method == 'POST':
        # Get text input from the form
        text = request.form['text']

        # Generate image from text
        #image = text_to_image(text)
        prompt = text
        negative_prompt = "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"

        image = pipe(
            prompt, 
            negative_prompt=negative_prompt,
            width=832,
            height=1216, 
            guidance_scale=7,
            num_inference_steps=28
        ).images[0]

        image.save("./output/asuka_test.png")

        # Save the generated image (for simplicity, you can send it directly to the user)
        # You can save it to a temporary file or memory depending on your needs
        #image.save("static/generated_image.png")

        return render_template('index.html', generated=True)

if __name__ == '__main__':
    app.run(debug=False)