Flask向HTML传递生成图片失败问题排查求助

你的代码核心问题出在HTML模板的条件判断，同时存在冗余代码，以下是具体问题和修复步骤：

主要问题点 #

1. HTML模板判断逻辑错误：模板中用{% if prediction %}判断是否显示图片，但Flask的predict函数返回的变量是img_data，这个判断永远不会触发，导致图片无法渲染。
2. 表单冗余配置：表单设置了enctype="multipart/form-data"，但你的逻辑里完全没处理上传的imagefile，属于多余配置。
3. HTML冗余代码：图片下方的<p>标签代码错误且无意义，还存在标签未闭合问题。

修复步骤 #

1. 修正HTML模板 #

• 将判断条件改为{% if img_data %}，确保Flask返回图片数据时能正确渲染。
• 移除无用的文件输入框和多余<p>标签。
• 移除表单的enctype="multipart/form-data"（无需处理文件上传）。

2. Flask代码微调（可选优化） #

• 添加异常捕获，避免生成图片失败时直接崩溃。
• 图片编码逻辑本身正确，无需修改。

修改后的完整代码 #

Flask代码 #

import os
import io
import base64
from PIL import Image
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
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(device)

@app.route('/', methods=['GET'])
def page_template():
    return render_template('index.html')

@app.route('/', methods=['POST'])
def predict():
    prompt_text = request.form['word']
    try:
        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'))
    except Exception as e:
        return render_template("index.html", error=str(e))

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

HTML模板（templates/index.html） #

<!DOCTYPE html>
<html>
<head>
    <title>Stable Diffusion 图片生成</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">
        <label for="promptfield" class="form-label">输入生成提示词</label>
        <input type="text" class="form-control" id="promptfield" name="word" placeholder="例如：一只在太空里的猫" required />
        <input class="btn btn-primary mt-3" type="submit" value="生成图片" >
    </form>

    {% if img_data %}
        <div class="text-center mt-4">
            <img id="picture" src="data:image/jpeg;base64,{{ img_data }}" class="img-fluid" style="max-width: 512px;">
        </div>
    {% endif %}

    {% if error %}
        <div class="text-center text-danger mt-3">
            <p>生成失败：{{ error }}</p>
        </div>
    {% endif %}
</body>
</html>

验证说明 #

1. 启动Flask应用后，访问首页输入提示词，点击生成按钮。
2. 生成的图片会以base64格式嵌入HTML，直接在页面中显示。
3. 如果生成过程出错，页面会显示错误信息，方便排查问题。

内容的提问来源于stack exchange，提问作者bullor