flask 前后台文件多张图片api；streamlit、gradio多图片页面展示

1、flask 前后台文件多张图片api

send_file 传递zip：
send_file(zip_data, mimetype=‘application/zip’, as_attachment=True, download_name=‘images.zip’)

from flask import Flask, Response, request,send_file
from PIL import Image
import torch
import io
from diffusers import PixArtAlphaPipeline
import zipfile

app = Flask(__name__)



## 初始化模型,第一次调用/image会来执行一次
@app.before_request
def load_model():
    if not hasattr(app, 'pipe'):
        torch.cuda.set_device(0)
        app.pipe = PixArtAlphaPipeline.from_pretrained("/ai/PixArt-XL-2-1024-MS", torch_dtype=torch.float16).to("cuda")
        #app.pipe.enable_model_cpu_offload()

@app.route('/image', methods=['POST'])
def get_image():
    prompt = request.form.get('prompt')
    
    # 生成图像文件并创建PIL.Image.Image对象
    images = app.pipe(prompt, num_inference_steps=10,num_images_per_prompt=4).images
# 在内存中创建一个BytesIO对象
    zip_data = io.BytesIO()

    # 将多张图片保存到BytesIO对象中
    with zipfile.ZipFile(zip_data, mode='w') as zipf:
        for i, image in enumerate(images):
            image_data = io.BytesIO()
            image.save(image_data, format='PNG')
            image_data.seek(0)
            zipf.writestr(f'image_{i+1}.png', image_data.getvalue())

    # 将BytesIO对象的指针重置到开头
    zip_data.seek(0)

    # 返回BytesIO对象作为文件给前端下载
    return send_file(zip_data, mimetype='application/zip', as_attachment=True, download_name='images.zip')





if __name__ == '__main__':
    app.run(host='192.1***', port=7889)

2、streamlit、gradio多图片展示

streamlit

st.columns(2)

import streamlit as st
from PIL import Image
import io
import pandas as pd
import numpy as np
import requests
from zipfile import ZipFile
import json

def translate(content):
    url = 'http://192***:11434/api/chat'
    headers = {'Content-Type': 'application/json'}
    

    pload = {
        "model": "qwen:1.8b-chat",
        "messages": [
            {"role": "system", "content": ""},
            {
                "role": "user",
                "content": f'''你是一个中英文翻译助手;
                后面双引号内中文问题请翻译成英文回复，问题是"{content}"'''
            }
                ],
        "stream": False
    }
    
    response = requests.post(url, headers=headers, json=pload)
    llm_text = json.loads(response.content)["message"]["content"]
    print("AI:",llm_text)
    return llm_text

# translate("大家好")

def get_image(prompt):
    
    # 发起 POST 请求并传递提示信息
    response = requests.post('http://192.1***:7889/image', data={'prompt': prompt})


    images = []
    if response.status_code == 200:
        # 解压缩接收到的ZIP文件
        with ZipFile(io.BytesIO(response.content), 'r') as zip_file:
            # 遍历ZIP文件中的所有文件
            for filename in zip_file.namelist():
                # 检查文件是否为图片文件（这里假设只有PNG格式的图片）
                if filename.endswith('.png'):
                    # 读取图片数据
                    image_data = zip_file.read(filename)
                    # 转换为图像对象
                    image = Image.open(io.BytesIO(image_data))
                    images.append(image)
    

    return images


# 初始化变量
if 'image1' not in st.session_state:
    st.session_state.image1 = None
if 'prompt' not in st.session_state:
    st.session_state.prompt = None
if 'prompt1' not in st.session_state:
    st.session_state.prompt1 = None
 

if __name__ == '__main__':
    st.title('AI文生图')

    # 创建文本输入框
    prompt = st.text_input('请输入提示词')
    # 将 prompt 存储在 st.session_state 中
    st.session_state.prompt1 = prompt    
    st.session_state.prompt = translate(prompt)
    
    # 监听按钮点击事件
    if st.button("生成图像"):
        with st.spinner("正在生成图像..."):
            # 生成图像
            st.session_state.image1 = get_image(st.session_state.prompt)
        st.text(st.session_state.prompt1)

        col1, col2 = st.columns(2)

        # 在第一排展示前两张图像
        with col1:
            st.image(st.session_state.image1[0], use_column_width=True)
            st.image(st.session_state.image1[1], use_column_width=True)

        # 在第二排展示后两张图像
        with col2:
            st.image(st.session_state.image1[2], use_column_width=True)
            st.image(st.session_state.image1[3], use_column_width=True)

gradio

gr.Gallery()

import gradio as gr
from PIL import Image
import io
import requests
import zipfile
import numpy as np


def get_images(prompt):
    response = requests.post('http://192****:7889/image', data={'prompt': prompt})
    images = []
    if response.status_code == 200:
        # 解压缩接收到的ZIP文件
        with zipfile.ZipFile(io.BytesIO(response.content), 'r') as zip_file:
            # 遍历ZIP文件中的所有文件
            for filename in zip_file.namelist():
                # 检查文件是否为图片文件（这里假设只有PNG格式的图片）
                if filename.endswith('.png'):
                    # 读取图片数据
                    image_data = zip_file.read(filename)
                    # 转换为图像对象
                    image = Image.open(io.BytesIO(image_data))
                    images.append(image)


    return images

def generate_images(prompt):
    images = get_images(prompt)
    return images

iface = gr.Interface(fn=generate_images, inputs="text", outputs=gr.Gallery(
            label="生成图片",height='auto',columns=2), title="AI文生图")
iface.launch(server_name="192.****2")