如何用Python实现Ollama+Llama2+Stable Diffusion的多任务生成（自动区分任务）

核心逻辑是让Llama 2先对用户请求做结构化解析，明确判断当前任务是文本生成还是图像生成，再根据判断结果调用对应的工具链。关键在于给LLM明确的指令，让它输出机器可解析的结构化结果（比如JSON），避免歧义。

1. 安装依赖 #

先安装所需的Python包：

pip install langchain-ollama diffusers transformers torch accelerate pillow

2. 配置LLM的任务识别提示 #

给Llama 2传入明确的指令，强制它输出固定格式的JSON，示例提示词：

你现在是任务分类器，需要判断用户的请求属于文本生成任务还是图像生成任务。请严格按照以下格式输出结果，不要添加任何多余内容：
{"task_type": "text" | "image", "content": "用户请求的核心内容，比如生成一段科幻故事、生成一只蓝色的猫的图片"}

3. 编写任务解析函数 #

用LangChain调用Ollama托管的Llama 2，完成任务类型判断：

import json
from langchain_ollama import ChatOllama

def parse_task(user_query: str) -> dict:
    llm = ChatOllama(model="llama2", temperature=0)
    prompt = f"""
你现在是任务分类器，需要判断用户的请求属于文本生成任务还是图像生成任务。请严格按照以下格式输出结果，不要添加任何多余内容：
{{"task_type": "text"|"image", "content": "用户请求的核心内容"}}

用户请求：{user_query}
"""
    response = llm.invoke(prompt).content
    # 处理可能的格式问题，比如LLM输出的多余换行或空格
    try:
        return json.loads(response.strip())
    except json.JSONDecodeError:
        # 解析失败时默认返回文本生成任务
        return {"task_type": "text", "content": user_query}

4. 实现文本生成与图像生成工具 #

文本生成函数 #

直接调用Ollama的Llama 2完成文本生成：

def generate_text(content: str) -> str:
    llm = ChatOllama(model="llama2", temperature=0.7)
    response = llm.invoke(f"根据以下内容生成文本：{content}").content
    return response

图像生成函数 #

用Hugging Face Diffusers调用Stable Diffusion（注意：需要本地有GPU加速，或者使用云端推理）：

from diffusers import StableDiffusionPipeline
import torch

def generate_image(content: str, output_path: str = "generated_image.png") -> str:
    # 加载Stable Diffusion模型，这里用默认的v1-5，也可以替换成其他版本
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float16
    ).to("cuda")  # 如果没有GPU，改成"cpu"，但速度会很慢
    
    image = pipe(content).images[0]
    image.save(output_path)
    return f"图像已保存至：{output_path}"

5. 主调度逻辑 #

整合所有模块，实现自动判断与执行：

def main():
    user_query = input("请输入你的请求：")
    task_info = parse_task(user_query)
    
    if task_info["task_type"] == "image":
        result = generate_image(task_info["content"])
    else:
        result = generate_text(task_info["content"])
    
    print("\n执行结果：")
    print(result)

if __name__ == "__main__":
    main()

注意事项 #

• 确保Ollama已安装并启动Llama 2模型（执行ollama pull llama2）
• Stable Diffusion首次运行会自动下载模型，需要充足的磁盘空间（约4GB）
• 如果没有GPU，建议使用Stable Diffusion的云端推理服务替代本地部署
• 可优化任务解析提示词，处理模糊请求（比如用户说“帮我创作森林内容”，可默认判断为文本生成或增加二次确认逻辑）

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