I would like to send image with Flask app when I click to predict button in my html code.
Code is supposed to get prompt and run ML model and generate image based on input prompt. I can generate the picture but somehow I can not send it to html or html doesn't show the image.
Anyone can check my flask and html code where the problem is ? My flask image send method is wrong or html part is wrong ? How to show the image in html ?
import os
import io
import base64
from PIL import Image
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
import transformers
from flask import Flask, render_template, request
app = Flask(__name__)
lms = LMSDiscreteScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear"
)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
pipe = StableDiffusionPipeline.from_pretrained(
"CompVis/stable-diffusion-v1-4",
scheduler=lms,
use_auth_token=True,
cache_dir=os.getenv("cache_dir", "./models")
).to("cuda")
@app.route('/', methods=['GET'])
def page_template():
return render_template('index.html')
@app.route('/', methods=['POST'])
def predict():
prompt_text = request.form['word']
with autocast("cuda"):
image = pipe(prompt_text)["sample"][0]
data = io.BytesIO()
image.save(data, "JPEG")
encoded_img_data = base64.b64encode(data.getvalue())
return render_template("index.html", img_data=encoded_img_data.decode('utf-8'))
if __name__ == '__main__':
app.run()
<!DOCTYPE html>
<html>
<head>
<title>Tutorial</title>
<link href="https://cdn.jsdelivr.net/npm/bootstrap@5.0.2/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-EVSTQN3/azprG1Anm3QDgpJLIm9Nao0Yz1ztcQTwFspd3yD65VohhpuuCOmLASjC" crossorigin="anonymous">
</head>
<body>
<h1 class="text-center">Stable Diffusion</h1>
<form class="p-3 text-center" action='/', method="post" enctype="multipart/form-data" >
<label for="promptfield" class="form-label">StableDiffusionPrompt</label>
<input type="text" class="form-control" id="promptfield" name="word" placeholder="Please enter your prompt" />
<input class="form-control" type="file" name="imagefile" >
<input class="btn btn-primary mt-3" type="submit" value="Predict Image" >
</form>
{% if prediction %}
<img id="picture" src="data:image/jpeg;base64,{{ img_data }}">
<p> img id="picture" src="data:image/jpeg;base64,{{ img_data }}"</p>
{% endif %}
</body>
</html>