YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
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:
- Analyze it with a vision model (Qwen2-VL)
- Generate embeddings
- Store in the vector database
- 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
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.
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
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):
| Variable | Default | Description |
|---|---|---|
HF_API_TOKEN |
(required) | Hugging Face API token |
DATABASE_URL |
sqlite+aiosqlite:///./storage/app.db |
Database connection string |
STORAGE_DIR |
./storage/images |
Image storage directory |
CHROMA_PERSIST_DIR |
./storage/chroma |
ChromaDB persistence directory |
VISION_MODEL |
Qwen/Qwen2-VL-7B-Instruct |
Vision-language model |
EMBEDDING_MODEL |
sentence-transformers/all-MiniLM-L6-v2 |
Embedding model |
GENERATION_MODEL |
black-forest-labs/FLUX.1-schnell |
Image generation model |
🔧 Models Used
| Component | Model | Purpose |
|---|---|---|
| Vision | Qwen/Qwen2-VL-7B-Instruct | Image → Text description |
| Embedding | sentence-transformers/all-MiniLM-L6-v2 | Text → 384-dim vector |
| Generation | black-forest-labs/FLUX.1-schnell | Text → Image |
| Fallback Gen | Pollinations AI | Free image generation (no API key) |
| Fallback Vision | Salesforce/blip-image-captioning-large | Basic image captioning |
📝 License
This project is for educational and development purposes.