This project provides an API wrapper around a Stable Diffusion model using FastAPI, enabling image generation via HTTP requests. The model is containerized using Docker for easy deployment.
- Uses PyTorch with CUDA acceleration for fast inference.
- Runs a FastAPI server to handle image generation requests.
- Generates images from text prompts and saves them inside the container.
- Fully containerized with Docker and supports GPU acceleration.
Ensure you have the following installed:
- Docker
- Docker Compose
- NVIDIA GPU with CUDA support (for faster processing)
- NVIDIA Container Toolkit (setup guide)
git clone <your-repo-url>
cd <your-repo-folder>Run the following command to build and save the pretrained model to the folder 'sd2_model'
python model.pyRun the following command to build the Docker image:
docker-compose buildThen, start the server:
docker-compose upThis will launch a FastAPI server on http://localhost:8000.
Once the container is running, open another terminal and execute:
python test.pyThis script sends a request to the API to generate an image.
After starting the container, you can manually test the API using:
curl -X POST "http://localhost:8000/generate-image/" -H "Content-Type: application/json" -d '{"prompt": "A fantasy landscape at sunset"}'The generated image will be saved inside the output/ directory in your local directory.
To stop the running container, use:
docker-compose down- If you see an error related to CUDA, ensure that your GPU drivers and NVIDIA Container Toolkit are correctly installed.
- If the build process takes too long, ensure that your model files are properly copied or mounted using volumes in
docker-compose.yml. - For permission issues, try running
sudo docker-compose upif on Linux.
Enjoy generating images with Stable Diffusion! 🚀