【NVIDIA NIM 黑客松训练营】文生图小应用

项目简介

本项目演示了一个基于NIM平台的生成式AI模型构建的简易Demo。
Demo使用了meta / llama3-70b-instruct和nvidia / consistory两个模型，首先，通过优化meta / llama3-70b-instruct模型的默认英文输出，使之更加符合中文用户的交流习惯；其次，通过智能分析用户的输入，判断其是否有绘画意图，从而决定是否调用nvidia / consistory模型来生成相应的图像。
整个项目采用Python语言，并借助FastAPI框架搭建了服务端，同时使用HTML技术构建了Web前端界面，为用户提供了一个直观且易用的操作平台。

应用链接

项目结构

work/
│
├── images/ # 存储生成的图像
├── chat.py # 后端逻辑处理脚本
└── chat.html # 前端交互界面

技术方案与实现步骤

安装依赖

在开始之前，确保安装了必要的Python包，以支持FastAPI框架和OpenAI接口的调用。

pip install fastapi uvicorn openai jinja2 requests

创建聊天应用

1. 首先要在NIM平台获取你的API KEY替换代码中的相应位置，以便能够顺利调用API。

from openai import OpenAI
from fastapi import FastAPI, Query
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
import requests, base64
from datetime import datetime
import json

host="127.0.0.1"
port=8080

nim_api_key = "你的apiKey"

invoke_url = "https://ai.api.nvidia.com/v1/genai/nvidia/consistory"
headers = {
   
    "Authorization": "Bearer " + nim_api_key,
    "Accept": "application/json",
}
client = OpenAI(
  base_url = "https://integrate.api.nvidia.com/v1",
  api_key = nim_api_key
)

2. 配置静态资源访问路径，配置跨域访问,供下面生成图片使用

app = FastAPI()
app.mount("/static", StaticFiles(directory="D:/WorkSpace/Python/nim-text2img/imgs"), name="static")
templates = Jinja2Templates(directory="templates")

@app.get("/static", response_class=HTMLResponse)
async def read_root(request: Request):
   return templates.TemplateResponse("chat.html", {
   "request": request})

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,  # 允许传递凭证（如 cookies）
    allow_methods=["*"],  # 允许所有 HTTP 方法
    allow_headers=["*"],  # 允许所有头部
)

3. 实现draw_image方法，接受四个传参，分别是图片主题prompt、主题词素、图片风格、场景描述，这四个参数是要传给nvidia / consistory模型的提示prompt，模型会根据这些参数进行绘画，如果想要一个好的图片质量，这四个参数尤为重要。方法返回的是图片链接。

def draw_image(subject_prompt, subject_tokens, style_prompt, scene_prompt1):
    print("subject_prompt=====",subject_prompt)
    print("subject_tokens=====",subject_tokens)
    print("style_prompt=====",style_prompt)
    print("scene_prompt1=====",scene_prompt1)
    payload = {
   
        "mode": 'init',
        "subject_prompt": subject_prompt,
        "subject_tokens": subject_tokens,
        "subject_seed": 43,
        "style_prompt": style_prompt,
        "scene_prompt1": scene_prompt1,
        "scene_prompt2": scene_prompt1,
        "negative_prompt": "",
        "cfg_scale": 5,
        "same_initial_noise": False
    }

    response = requests.post(invoke_url, headers=headers, json=payload)
    response.raise_for_status()

    data = response.json()
    current_time = datetime.now(