commit#15
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hello @hndrr, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
Summary of Changes
Hello! gemini-code-assist here, providing a summary of this pull request.
This PR focuses on integrating image generation capabilities into the FramePack-FastAPI application, which previously handled only video generation. It introduces new API endpoints for single image generation, batch image generation, and image style transfer (kisekaeichi). The changes also include updates to the job queue system to handle different job types, a new worker specifically for image tasks, and refinements to the Docker setup to better support both a lightweight development mode (without large models) and a full GPU-accelerated inference mode.
The PR also adds initial tests for the new image API endpoints and includes documentation outlining the integration plan and revised Docker usage.
Highlights
- Image Generation API: Adds new API endpoints:
/api/generate-imagefor single image generation,/api/batch-imagesfor generating multiple images from a list of prompts,/api/transfer-imagefor image style transfer, and/api/image/{job_id}to retrieve generated images. - Image Worker Implementation: Introduces a new worker module (
api/worker_image.py) containing the core logic for image generation, including prompt encoding, latent preparation, diffusion steps, and saving the final image with metadata. Placeholder functions for batch and transfer are included. - Queue System Enhancement: The job queue manager (
api/queue_manager.py) is updated to support different job types (video, image, batch_image, image_transfer) and store job-specific data and results, providing a unified system for both video and image tasks. - Docker Configuration Refinement: Docker files and scripts are updated, along with new documentation (
README_DOCKER_REVISED.md), to clarify and support distinct Docker environments: a lightweight development server (API testing without models) and a GPU-accelerated inference server (for actual generation). - Development Mode: A mechanism is added to the API (
api/api.py) to skip loading large models based on theMODELS_DISABLEDenvironment variable, enabling a faster startup and lower resource usage for development and testing. - API Tests: Adds initial pytest tests (
tests/test_image_api.py) for the new image generation endpoints to verify basic functionality and request validation.
Changelog
Click here to see the changelog
- .dockerignore
- Added comprehensive ignore patterns for Git, Python, virtual environments, IDE files, OS files, project-specific outputs (outputs/, temp_queue_images/, archive/, hf_download/, job_queue.json, quick_prompts.json, *.log), Docker files, documentation, cache directories, and temporary files.
- .gitignore
- Added
CLAUDE.mdand.mcp/to the ignore list.
- Added
- Dockerfile
- Added a new Dockerfile based on
nvidia/cuda:12.1-devel-ubuntu22.04. - Installs Python 3.10, pip, curl, git, build essentials, and various libraries (libgl1-mesa-glx, libglib2.0-0, libsm6, libxext6, libxrender-dev, libgomp1, ffmpeg).
- Sets up a symlink for
pythontopython3.10. - Copies and uses
uvfor dependency management. - Sets the working directory to
/app. - Copies
pyproject.tomlandrequirements.txt. - Creates a virtual environment using
uv venvand sets the PATH. - Installs Python dependencies using
uv pip install -r requirements.txt. - Copies the application code.
- Creates necessary directories:
outputs/images,temp_queue_images,loras. - Creates a model mount point
/app/hf_downloadand declares it as a VOLUME. - Sets environment variables
PYTHONPATH,PYTHONUNBUFFERED,HF_HOME,TRANSFORMERS_CACHE,HF_DATASETS_CACHE. - Exposes port 8000.
- Adds a health check endpoint (
curl -f http://localhost:8000/docs). - Sets the default command to run uvicorn (
uvicorn api.api:app --host 0.0.0.0 --port 8000).
- Added a new Dockerfile based on
- README_DOCKER.md
- Added a new Japanese language README providing a guide for running FramePack-FastAPI with Docker.
- Describes available Dockerfiles (CPU/GPU) and Docker Compose files (CPU/GPU).
- Provides quick start, manual build, and run instructions for both CPU and GPU versions.
- Details customization options (port, environment variables, volumes).
- Outlines system requirements for CPU and GPU.
- Includes troubleshooting tips for common issues (GPU recognition, memory, model download, port conflicts).
- Explains API access URLs.
- Provides instructions for updating, backing up data, and stopping/removing containers.
- README_DOCKER_REVISED.md
- Added a new Japanese language README providing a revised and clarified guide for Docker usage.
- Explicitly states that the CPU version (
Dockerfile) is for development/testing only and not realistic for inference. - Details the purpose and requirements for the Development Server (
Dockerfile) and GPU Inference Server (Dockerfile.gpu). - Provides usage examples for development/testing and full GPU operation.
- Outlines realistic operational workflows (personal, team, production).
- Provides detailed system requirements for both server types.
- Details Docker configuration including port settings (8000 for dev, 8001 for GPU), volumes, and environment variables.
- Includes troubleshooting specific to the clarified setup.
- Suggests recommended workflows and cost optimization tips.
- Lists important caveats regarding CPU inference, GPU requirements, and storage.
- api/api.py
- Imported
datetimeanduuid(line 20-21). - Imported
worker_imagemodule (line 30). - Modified the
lifespanfunction to conditionally load models based on theMODELS_DISABLEDenvironment variable (line 51-65). - Added Pydantic models for image generation requests (
ImageGenerationRequest,BatchImageRequest,ImageTransferRequest) and response (ImageGenerationResponse) (line 189-227). - Added new API endpoint
POST /api/generate-imagefor single image generation (line 759-792). - Added new API endpoint
POST /api/batch-imagesfor batch image generation (line 795-828). - Added new API endpoint
POST /api/transfer-imagefor image style transfer (line 831-864). - Added new API endpoint
GET /api/image/{job_id}to download generated image files (line 867-882).
- Imported
- api/image_models.py
- Added a new file
api/image_models.py. - Defines
IMAGE_GENERATION_CONFIGdictionary with image-specific parameters (line 8-17). - Includes
get_rope_configfunction to get RoPE configuration based on dimensions (line 20-42). - Includes
get_image_latent_shapefunction to calculate latent tensor shape (line 45-72). - Includes
validate_image_dimensionsfunction to adjust dimensions for VAE compatibility (line 75-97). - Includes
ImageGenerationHelperclass with static methods for preparing image metadata and estimating memory requirements (line 101-191).
- Added a new file
- api/queue_manager.py
- Imported
Anyfor type hinting (line 8). - Added
job_type,data, andresultfields to theQueuedJobdataclass (line 55-59). - Updated
to_dictmethod to includejob_type,data, andresult(line 91-93). - Updated
from_dictmethod to loadjob_type,data, andresult(line 146-148). - Added
add_jobfunction to add a generic job (including image jobs) to the queue (line 284-314). - Added
update_job_resultfunction to update the result field of a job and save the queue (line 317-353). - Modified
get_queue_statusto includejob_typein the status output (line 662).
- Imported
- api/worker_image.py
- Added a new file
api/worker_image.py. - Imports necessary libraries and helper modules for image generation (line 1-23).
- Includes
encode_cropped_promptfunction for processing prompts with dual text encoders (line 26-107). - Includes
prepare_latentsfunction for initializing diffusion latents (line 110-118). - Implements
process_image_generationfunction for single image generation logic (line 121-325). - Includes placeholder
process_batch_imagesandprocess_image_transferfunctions (line 327-342).
- Added a new file
- docker-build.sh
- Modified the script to build the
Dockerfile(likely the GPU version based on content). - Added
--pushand--tagoptions. - Prints build configuration and instructions for running the resulting image.
- Modified the script to build the
- docker-compose.yml
- Modified the compose file to define a service
framepack-apiusing theDockerfile. - Maps port 8000.
- Sets environment variables including
NVIDIA_VISIBLE_DEVICES,HF_HOME, etc. - Mounts volumes for outputs, temp images, loras, job queue, and hf_download.
- Configures deployment resources to require 1 NVIDIA GPU.
- Includes a health check.
- Modified the compose file to define a service
- docker-run.sh
- Modified the script to use
docker-compose -f docker-compose.yml up. - Added
--detachedand--portoptions. - Creates necessary directories (
outputs/images,temp_queue_images,loras,hf_download). - Temporarily modifies the compose file to change the port mapping if requested.
- Provides instructions for viewing logs and stopping containers when run in detached mode.
- Modified the script to use
- integration_plan.md
- Added a new Japanese language document outlining the plan for integrating image generation features.
- Details new API endpoints, file structure changes, and modifications to existing files.
- Defines implementation phases (single image, batch, transfer, optimization).
- Discusses technical considerations (memory, concurrency, compatibility).
- Outlines expected outcomes, risks, and mitigation strategies.
- Lists next steps for the project.
- pyproject.toml
- Updated the project description to include 'with image and video generation' (line 8).
- Added
httpx,mypy,pytest,pytest-asyncio, andruffto thedevoptional dependencies (line 37-42).
- start_server.sh
- Modified the script to use
uvfor virtual environment management and dependency installation. - Sets
PYTHONPATH. - Runs uvicorn via
uv runon port 8080 with--reload.
- Modified the script to use
- tests/test_image_api.py
- Added a new file
tests/test_image_api.py. - Contains pytest tests for the new image generation API endpoints (
/api/generate-image,/api/batch-images,/api/transfer-image,/api/image/{job_id}). - Includes tests for endpoint existence, minimal/full parameters, batch requests, image transfer requests, image download, and validation errors.
- Added a new file
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces image generation capabilities to the FramePack-FastAPI project, adding endpoints for single image generation, batch processing, and image transfer. It also includes updates to the Docker configuration and documentation to support these new features.
The core implementation for single image generation (worker_image.py) seems to follow a standard diffusion process. The updates to the queue_manager to handle different job types are well-integrated.
However, there are some inconsistencies in the Docker setup scripts and compose file regarding which Dockerfile (Dockerfile vs Dockerfile.gpu) is used for which purpose, which could lead to confusion. Additionally, while basic API endpoint tests are included, there is a lack of functional tests for the actual image generation logic.
Overall, this is a significant step towards a more versatile API, but addressing the points below will improve clarity, maintainability, and confidence in the new features.
Summary of Findings
- Docker Configuration Inconsistency: There is a significant inconsistency in the Docker setup. The
Dockerfileis CUDA-based but described as a CPU/dev version in the READMEs. Thedocker-compose.ymlfile uses thisDockerfilebut requests GPU resources, contradicting the READMEs which state this compose file is for the CPU version. Thedocker-build.shanddocker-run.shscripts also reference theDockerfilewhile mentioning GPU requirements. This needs to be aligned so thatDockerfileis a true CPU/dev image (without CUDA base or GPU resource requests) andDockerfile.gpu(or similar) is used for the GPU version, with corresponding updates to the scripts and compose files, matching the documentation inREADME_DOCKER_REVISED.md. - Lack of Functional Tests for Image Generation: The new test file
tests/test_image_api.pyincludes tests for endpoint existence and request validation, which is good. However, it lacks functional tests that verify the actual image generation process, batch processing, or image transfer. Tests should be added to submit a job, wait for its completion, and verify that an output file is generated and downloadable. - Hardcoded Paths and Settings: Some paths and settings (like the image output directory and thumbnail size) are hardcoded within
worker_image.py. These should ideally be read from thesettingsmodule for better maintainability and configurability. - Placeholder Implementations (TODOs): The
process_batch_imagesandprocess_image_transferfunctions inworker_image.pyare currently marked as TODOs. While the API endpoints are present, the core logic for these features is missing. uuidImport Scope: Theuuidmodule is imported inside the image generation endpoint functions inapi/api.py. It should be imported once at the top of the file.- In-place
sedindocker-run.sh: Thedocker-run.shscript modifies thedocker-compose.ymlfile in place usingsedto change the port. Using environment variables directly in the compose file is a more standard and less intrusive approach. - uv.lock file: The
uv.lockfile is included. This is a minor point related to dependency management workflow choices and was not commented on in the code review due to the configured severity settings.
Merge Readiness
This pull request introduces significant new functionality. However, due to the high-severity issues identified, particularly the inconsistencies in the Docker configuration which could prevent correct setup and usage, and the lack of functional tests for the core image generation logic, I recommend requesting changes before merging. Addressing these issues will ensure the new features are correctly implemented, maintainable, and reliable. Please have other reviewers assess the changes once the feedback is addressed. I am unable to approve this pull request.
|
|
||
| # Set image name and dockerfile | ||
| IMAGE_NAME="framepack-fastapi" | ||
| DOCKERFILE="Dockerfile" |
There was a problem hiding this comment.
This script references Dockerfile but the warning below mentions GPU requirements. Given Dockerfile is described as the CPU/dev version in the READMEs, this script should probably reference Dockerfile.gpu if it's intended for building the GPU image, or be updated to handle both.
| 1. **Dockerfile** - CPU版（軽量、推論のみ） | ||
| 2. **Dockerfile.gpu** - GPU版（CUDA対応、高性能） |
There was a problem hiding this comment.
This section correctly identifies Dockerfile as CPU and Dockerfile.gpu as GPU. This aligns with the intent described in README_DOCKER_REVISED.md, but contradicts the actual Dockerfile content and the usage in docker-compose.yml and the run scripts (see other comments). Please ensure consistency.
| 1. **docker-compose.yml** - CPU版用 | ||
| 2. **docker-compose.gpu.yml** - GPU版用 |
There was a problem hiding this comment.
Similarly, this section correctly identifies docker-compose.yml for CPU and docker-compose.gpu.yml for GPU. This aligns with README_DOCKER_REVISED.md but contradicts the current docker-compose.yml which requests GPU resources.
| """ | ||
| Test cases for image generation API endpoints | ||
| """ | ||
| import pytest | ||
| import json | ||
| import base64 | ||
| import io | ||
| from PIL import Image | ||
| from fastapi.testclient import TestClient | ||
| from api.api import app | ||
|
|
||
| client = TestClient(app) | ||
|
|
||
|
|
||
| def create_test_image_base64(width: int = 512, height: int = 512) -> str: | ||
| """Create a test image encoded as base64""" | ||
| image = Image.new('RGB', (width, height), color='red') | ||
| buffer = io.BytesIO() | ||
| image.save(buffer, format='PNG') | ||
| image_data = base64.b64encode(buffer.getvalue()).decode('utf-8') | ||
| return f"data:image/png;base64,{image_data}" | ||
|
|
||
|
|
||
| class TestImageGenerationAPI: | ||
| """Test class for image generation endpoints""" | ||
|
|
||
| def test_generate_image_endpoint_exists(self): | ||
| """Test that the image generation endpoint exists""" | ||
| response = client.post("/api/generate-image", json={ | ||
| "prompt": "A beautiful landscape" | ||
| }) | ||
| # Should not return 404 (endpoint exists) | ||
| assert response.status_code != 404 | ||
|
|
||
| def test_generate_image_with_minimal_params(self): | ||
| """Test image generation with minimal parameters""" | ||
| request_data = { | ||
| "prompt": "A cute cat sitting on a chair" | ||
| } | ||
|
|
||
| response = client.post("/api/generate-image", json=request_data) | ||
|
|
||
| # Should accept the request (return job_id) | ||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| assert len(data["job_id"]) > 0 | ||
| else: | ||
| # If models aren't loaded, expect 503 | ||
| assert response.status_code == 503 | ||
|
|
||
| def test_generate_image_with_full_params(self): | ||
| """Test image generation with all parameters""" | ||
| request_data = { | ||
| "prompt": "A detailed fantasy landscape with mountains and rivers", | ||
| "negative_prompt": "blurry, low quality, distorted", | ||
| "seed": 42, | ||
| "steps": 25, | ||
| "cfg": 7.5, | ||
| "width": 768, | ||
| "height": 512, | ||
| "lora_paths": [], | ||
| "lora_scales": [] | ||
| } | ||
|
|
||
| response = client.post("/api/generate-image", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 # Models not loaded | ||
|
|
||
| def test_batch_images_endpoint(self): | ||
| """Test batch image generation endpoint""" | ||
| request_data = { | ||
| "prompts": [ | ||
| "A red apple", | ||
| "A blue car", | ||
| "A green tree" | ||
| ], | ||
| "batch_size": 3, | ||
| "steps": 20 | ||
| } | ||
|
|
||
| response = client.post("/api/batch-images", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 | ||
|
|
||
| def test_image_transfer_endpoint(self): | ||
| """Test image transfer endpoint""" | ||
| source_image = create_test_image_base64(256, 256) | ||
| target_image = create_test_image_base64(256, 256) | ||
|
|
||
| request_data = { | ||
| "source_image": source_image, | ||
| "target_image": target_image, | ||
| "prompt": "Transfer the style from source to target", | ||
| "transfer_strength": 0.8 | ||
| } | ||
|
|
||
| response = client.post("/api/transfer-image", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 | ||
|
|
||
| def test_image_download_endpoint(self): | ||
| """Test image download endpoint""" | ||
| # Test with a non-existent job ID | ||
| response = client.get("/api/image/nonexistent123") | ||
| assert response.status_code == 404 | ||
|
|
||
| def test_invalid_image_generation_params(self): | ||
| """Test image generation with invalid parameters""" | ||
| # Test with empty prompt | ||
| response = client.post("/api/generate-image", json={ | ||
| "prompt": "" | ||
| }) | ||
| # Should either accept empty prompt or return validation error | ||
| assert response.status_code in [200, 422, 503] | ||
|
|
||
| # Test with negative steps | ||
| response = client.post("/api/generate-image", json={ | ||
| "prompt": "test", | ||
| "steps": -5 | ||
| }) | ||
| # Should return validation error | ||
| assert response.status_code == 422 | ||
|
|
||
| # Test with invalid CFG | ||
| response = client.post("/api/generate-image", json={ | ||
| "prompt": "test", | ||
| "cfg": -1.0 | ||
| }) | ||
| # Should return validation error | ||
| assert response.status_code == 422 | ||
|
|
||
| def test_batch_images_validation(self): | ||
| """Test batch image generation validation""" | ||
| # Test with empty prompts list | ||
| response = client.post("/api/batch-images", json={ | ||
| "prompts": [] | ||
| }) | ||
| # Should return validation error | ||
| assert response.status_code == 422 | ||
|
|
||
| # Test with too large batch size | ||
| response = client.post("/api/batch-images", json={ | ||
| "prompts": ["test1", "test2"], | ||
| "batch_size": 20 # Exceeds max of 8 | ||
| }) | ||
| # Should return validation error | ||
| assert response.status_code == 422 | ||
|
|
||
| def test_image_transfer_validation(self): | ||
| """Test image transfer validation""" | ||
| # Test with invalid base64 image | ||
| response = client.post("/api/transfer-image", json={ | ||
| "source_image": "invalid_base64", | ||
| "target_image": "invalid_base64", | ||
| "prompt": "test" | ||
| }) | ||
| # Should either return validation error or process with error | ||
| assert response.status_code in [422, 503] | ||
|
|
||
| # Test with invalid transfer strength | ||
| source_image = create_test_image_base64() | ||
| target_image = create_test_image_base64() | ||
|
|
||
| response = client.post("/api/transfer-image", json={ | ||
| "source_image": source_image, | ||
| "target_image": target_image, | ||
| "prompt": "test", | ||
| "transfer_strength": 2.0 # Should be between 0.0 and 1.0 | ||
| }) | ||
| # Should return validation error | ||
| assert response.status_code == 422 | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| pytest.main([__file__, "-v"]) No newline at end of file |
There was a problem hiding this comment.
While the tests cover endpoint existence and basic request validation, they do not test the core functionality of image generation, batch processing, or image transfer. There are no tests that wait for a job to complete and then verify the output (e.g., download the generated image). Adding functional tests that cover the full job lifecycle would significantly improve confidence in the correctness of the new features.
| # Set compose file | ||
| COMPOSE_FILE="docker-compose.yml" | ||
| echo "🚀 Starting FramePack-FastAPI..." | ||
| echo "⚠️ This requires NVIDIA GPU with 24GB+ VRAM!" |
There was a problem hiding this comment.
This warning mentions GPU requirements, suggesting the script is for the GPU version, but it uses docker-compose.yml which is described as the CPU version in the READMEs. Please align the script's purpose and the compose file it uses with the documentation.
| def test_generate_image_with_minimal_params(self): | ||
| """Test image generation with minimal parameters""" | ||
| request_data = { | ||
| "prompt": "A cute cat sitting on a chair" | ||
| } | ||
|
|
||
| response = client.post("/api/generate-image", json=request_data) | ||
|
|
||
| # Should accept the request (return job_id) | ||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| assert len(data["job_id"]) > 0 | ||
| else: | ||
| # If models aren't loaded, expect 503 | ||
| assert response.status_code == 503 |
There was a problem hiding this comment.
This test checks if the /api/generate-image endpoint accepts a minimal request and returns a job ID (or 503 if models are not loaded). This is a basic check, but it doesn't verify if the job actually starts processing or completes successfully.
| def test_generate_image_with_full_params(self): | ||
| """Test image generation with all parameters""" | ||
| request_data = { | ||
| "prompt": "A detailed fantasy landscape with mountains and rivers", | ||
| "negative_prompt": "blurry, low quality, distorted", | ||
| "seed": 42, | ||
| "steps": 25, | ||
| "cfg": 7.5, | ||
| "width": 768, | ||
| "height": 512, | ||
| "lora_paths": [], | ||
| "lora_scales": [] | ||
| } | ||
|
|
||
| response = client.post("/api/generate-image", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 # Models not loaded |
There was a problem hiding this comment.
Similar to the minimal params test, this checks if the endpoint accepts a request with all parameters. It also doesn't verify the actual generation process.
| def test_batch_images_endpoint(self): | ||
| """Test batch image generation endpoint""" | ||
| request_data = { | ||
| "prompts": [ | ||
| "A red apple", | ||
| "A blue car", | ||
| "A green tree" | ||
| ], | ||
| "batch_size": 3, | ||
| "steps": 20 | ||
| } | ||
|
|
||
| response = client.post("/api/batch-images", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 |
There was a problem hiding this comment.
This test checks if the /api/batch-images endpoint accepts a request. Again, it doesn't verify the actual batch processing or completion.
| def test_image_transfer_endpoint(self): | ||
| """Test image transfer endpoint""" | ||
| source_image = create_test_image_base64(256, 256) | ||
| target_image = create_test_image_base64(256, 256) | ||
|
|
||
| request_data = { | ||
| "source_image": source_image, | ||
| "target_image": target_image, | ||
| "prompt": "Transfer the style from source to target", | ||
| "transfer_strength": 0.8 | ||
| } | ||
|
|
||
| response = client.post("/api/transfer-image", json=request_data) | ||
|
|
||
| if response.status_code == 200: | ||
| data = response.json() | ||
| assert "job_id" in data | ||
| assert "message" in data | ||
| else: | ||
| assert response.status_code == 503 |
There was a problem hiding this comment.
This test checks if the /api/transfer-image endpoint accepts a request with base64 images. It doesn't verify the actual transfer logic.
| # Load models (skip if in development mode) | ||
| if os.getenv("MODELS_DISABLED", "false").lower() == "true": | ||
| print("Models loading disabled - running in development mode") | ||
| loaded_models = {} | ||
| else: | ||
| try: | ||
| # Consider running blocking IO in a threadpool executor in async context | ||
| # e.g., await asyncio.to_thread(models.load_models) | ||
| # For simplicity now, keeping the direct call but be aware of potential blocking | ||
| loaded_models = models.load_models() | ||
| print("Models loaded successfully via lifespan.") | ||
| except Exception as e: | ||
| print(f"FATAL: Failed to load models on startup via lifespan: {e}") | ||
| traceback.print_exc() | ||
| loaded_models = {} |
There was a problem hiding this comment.
Adding logic to skip model loading based on the MODELS_DISABLED environment variable is a good way to support a lightweight development mode without large model downloads.
1 participant
No description provided.