image-rag-pipeline

Image-to-Image RAG Pipeline

A production-style AI image generation system with RAG (Retrieval-Augmented Generation) memory. Built with FastAPI, Hugging Face APIs, ChromaDB, and SQLite.

The system learns from stored user data and retrieves similar designs before generating new images, creating a context-aware image generation pipeline.

🏗️ Architecture

User Request (image/text)
        │
        ▼
┌─────────────────────┐
│   FastAPI Endpoints  │  /upload-image, /generate-image, /search-similar, /history
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Vision Service     │  Qwen2-VL via HF Inference API
│  (Image → Text)     │  Converts uploaded image to detailed description
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Embedding Service  │  sentence-transformers/all-MiniLM-L6-v2
│  (Text → Vector)    │  Generates 384-dim embeddings from descriptions
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Vector DB (ChromaDB)│  Store & retrieve similar image embeddings
│  RAG Retrieval       │  Top-K cosine similarity search
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Prompt Engineering │  Combines user input + vision output + RAG results
│  Service            │  into an optimized generation prompt
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Image Generation   │  FLUX.1-schnell via HF Inference API
│  Service            │  Generates new image from final prompt
└─────────┬───────────┘
          │
          ▼
┌─────────────────────┐
│  Storage Service    │  SQLite DB (metadata) + Local filesystem (images)
└─────────────────────┘

📁 Project Structure

image-rag-pipeline/
├── app/
│   ├── __init__.py
│   ├── main.py                  # FastAPI app entry point
│   ├── config.py                # Environment variables & settings
│   ├── api/
│   │   ├── __init__.py
│   │   └── routes.py            # All API endpoints
│   ├── services/
│   │   ├── __init__.py
│   │   ├── vision.py            # Qwen2-VL image understanding
│   │   ├── embedding.py         # Sentence-transformers embeddings
│   │   ├── generation.py        # FLUX.1 / SDXL image generation
│   │   └── prompt_engineer.py   # RAG-enhanced prompt building
│   ├── rag/
│   │   ├── __init__.py
│   │   └── vector_store.py      # ChromaDB vector operations
│   ├── models/
│   │   ├── __init__.py
│   │   ├── database.py          # SQLAlchemy ORM models
│   │   └── schemas.py           # Pydantic request/response schemas
│   └── utils/
│       ├── __init__.py
│       └── storage.py           # Local file storage
├── storage/
│   ├── images/                  # Uploaded & generated images
│   │   ├── uploaded/
│   │   └── generated/
│   ├── chroma/                  # ChromaDB persistence
│   └── app.db                   # SQLite database
├── requirements.txt
├── .env                         # Environment variables (create from .env.example)
├── .env.example                 # Template for env vars
└── README.md

🚀 Quick Start

1. Install Dependencies

cd image-rag-pipeline
pip install -r requirements.txt

2. Configure Environment

# Copy the example env file
cp .env.example .env

# Edit .env and add your Hugging Face API token
# Get your token from: https://huggingface.co/settings/tokens

Set your HF_API_TOKEN in the .env file:

HF_API_TOKEN=hf_your_actual_token_here

3. Run the Server

uvicorn app.main:app --reload --port 8000

The server will start at http://localhost:8000. Visit http://localhost:8000/docs for the interactive API documentation (Swagger UI).

📡 API Endpoints

POST /upload-image — Upload & Analyze Image

Upload an image to the RAG pipeline. The system will:

1. Analyze it with a vision model (Qwen2-VL)
2. Generate embeddings
3. Store in the vector database
4. Find similar images

Request:

curl -X POST http://localhost:8000/upload-image \
  -F "file=@my_photo.jpg"

Response:

{
  "image_id": "a1b2c3d4-...",
  "description": "A vibrant sunset over a mountain landscape with warm orange and purple tones...",
  "image_url": "/images/uploaded/a1b2c3d4_my_photo.jpg",
  "similar_images": [
    {
      "image_id": "e5f6g7h8-...",
      "description": "A mountain scene at dusk...",
      "image_url": "/images/uploaded/e5f6g7h8_sunset.png",
      "similarity_score": 0.8723,
      "source": "uploaded"
    }
  ],
  "message": "Image uploaded and indexed successfully. Found 1 similar images."
}

POST /generate-image — Generate Image with RAG

Generate a new image using the full RAG pipeline. The system retrieves similar images for context-aware generation.

Request:

curl -X POST http://localhost:8000/generate-image \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": "a futuristic city at sunset with flying cars",
    "style": "cyberpunk",
    "top_k": 5
  }'

Response:

{
  "image_id": "x9y0z1a2-...",
  "image_url": "/images/generated/x9y0z1a2.png",
  "user_prompt": "a futuristic city at sunset with flying cars",
  "final_prompt": "A futuristic cyberpunk scene of a futuristic city at sunset with flying cars, inspired by: ..., neon glowing lights, synthwave aesthetic, dark atmosphere",
  "similar_context": [...],
  "message": "Image generated successfully. Used 3 similar images as RAG context."
}

Available Styles: detailed, digital_art, photorealistic, anime, cyberpunk, watercolor, fantasy, 3d_render

POST /search-similar — Similarity Search

Search the vector store for similar images by text or by image ID.

By text:

curl -X POST http://localhost:8000/search-similar \
  -H "Content-Type: application/json" \
  -d '{"query": "sunset over mountains", "top_k": 5}'

By image ID:

curl -X POST http://localhost:8000/search-similar \
  -H "Content-Type: application/json" \
  -d '{"image_id": "a1b2c3d4-...", "top_k": 5}'

Response:

{
  "query": "sunset over mountains",
  "results": [
    {
      "image_id": "...",
      "description": "...",
      "image_url": "/images/...",
      "similarity_score": 0.9124,
      "source": "uploaded"
    }
  ],
  "total_results": 3
}

GET /history — Generation History

Retrieve paginated generation history.

curl "http://localhost:8000/history?page=1&page_size=10"

Response:

{
  "items": [
    {
      "id": "...",
      "user_input": "a futuristic city at sunset",
      "input_type": "text",
      "vision_output": null,
      "final_prompt": "...",
      "generated_image_url": "/images/generated/...",
      "similarity_score": 0.85,
      "created_at": "2025-01-15T10:30:00Z"
    }
  ],
  "total": 42,
  "page": 1,
  "page_size": 10
}

GET / — Health Check

curl http://localhost:8000/

GET /stats — System Statistics

curl http://localhost:8000/stats

🧠 How the RAG Pipeline Works

1. Upload Phase: When you upload an image, the vision model (Qwen2-VL) creates a detailed textual description. This description is embedded into a 384-dimensional vector and stored in ChromaDB alongside the image metadata.

2. Generation Phase: When you request image generation:

   • Your text prompt is converted to an embedding
   • ChromaDB finds the most similar images from your stored collection
   • The prompt engineer combines your input with RAG-retrieved context
   • The final optimized prompt is sent to the image generation model
   • The generated image and its embedding are stored back into the RAG system

3. Learning: Over time, as more images are uploaded and generated, the RAG system builds a richer knowledge base. Each new generation benefits from all previous interactions, making the system increasingly context-aware.

⚙️ Configuration

All settings are managed via environment variables (.env file):

🔧 Models Used

📝 License

This project is for educational and development purposes.