aquiles-image 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aquiles_image-0.1.0/PKG-INFO +47 -0
- aquiles_image-0.1.0/README.md +10 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/PKG-INFO +47 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/SOURCES.txt +24 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/dependency_links.txt +1 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/entry_points.txt +2 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/requires.txt +17 -0
- aquiles_image-0.1.0/aquiles_image.egg-info/top_level.txt +1 -0
- aquiles_image-0.1.0/aquilesimage/cli/__init__.py +1 -0
- aquiles_image-0.1.0/aquilesimage/cli/cli.py +147 -0
- aquiles_image-0.1.0/aquilesimage/configs/__init__.py +1 -0
- aquiles_image-0.1.0/aquilesimage/configs/configs.py +100 -0
- aquiles_image-0.1.0/aquilesimage/main.py +178 -0
- aquiles_image-0.1.0/aquilesimage/models/__init__.py +5 -0
- aquiles_image-0.1.0/aquilesimage/models/models.py +211 -0
- aquiles_image-0.1.0/aquilesimage/pipelines/__init__.py +1 -0
- aquiles_image-0.1.0/aquilesimage/pipelines/pipelines.py +255 -0
- aquiles_image-0.1.0/aquilesimage/runtime/__init__.py +2 -0
- aquiles_image-0.1.0/aquilesimage/runtime/requestscopedpipeline.py +289 -0
- aquiles_image-0.1.0/aquilesimage/runtime/scheduler.py +137 -0
- aquiles_image-0.1.0/aquilesimage/static/aquilesim.png +0 -0
- aquiles_image-0.1.0/aquilesimage/utils/__init__.py +1 -0
- aquiles_image-0.1.0/aquilesimage/utils/utils.py +94 -0
- aquiles_image-0.1.0/pyproject.toml +66 -0
- aquiles_image-0.1.0/setup.cfg +4 -0
- aquiles_image-0.1.0/test/test.py +11 -0
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: aquiles-image
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Aquiles-Image.
|
|
5
|
+
Author-email: Aquiles-ai / Fredy <riveraaai200678@gmail.com>
|
|
6
|
+
License: Apache License 2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/Aquiles-ai/Aquiles-Image
|
|
8
|
+
Project-URL: Issues, https://github.com/Aquiles-ai/Aquiles-Image/issues
|
|
9
|
+
Keywords: fastapi,ai
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
16
|
+
Classifier: Operating System :: OS Independent
|
|
17
|
+
Classifier: Environment :: Web Environment
|
|
18
|
+
Classifier: Topic :: Software Development :: Build Tools
|
|
19
|
+
Requires-Python: >=3.8
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
Requires-Dist: torch
|
|
22
|
+
Requires-Dist: torchvision
|
|
23
|
+
Requires-Dist: transformers
|
|
24
|
+
Requires-Dist: sentencepiece
|
|
25
|
+
Requires-Dist: fastapi
|
|
26
|
+
Requires-Dist: click>=8.0.0
|
|
27
|
+
Requires-Dist: uvicorn
|
|
28
|
+
Requires-Dist: ftfy
|
|
29
|
+
Requires-Dist: accelerate
|
|
30
|
+
Requires-Dist: xformers
|
|
31
|
+
Requires-Dist: protobuf
|
|
32
|
+
Provides-Extra: dev
|
|
33
|
+
Requires-Dist: pytest; extra == "dev"
|
|
34
|
+
Requires-Dist: black; extra == "dev"
|
|
35
|
+
Requires-Dist: isort; extra == "dev"
|
|
36
|
+
Requires-Dist: mypy; extra == "dev"
|
|
37
|
+
|
|
38
|
+
<h1 align="center">Aquiles-Image</h1>
|
|
39
|
+
|
|
40
|
+
<div align="center">
|
|
41
|
+
<img src="aquilesimage/static/aquilesim.png" alt="Aquiles-Image Logo" width="225"/>
|
|
42
|
+
</div>
|
|
43
|
+
|
|
44
|
+
<p align="center">
|
|
45
|
+
<strong>Easy, fast and cheap Diffusion Models that work for everyone.</strong><br/>
|
|
46
|
+
🚀 FastAPI • Diffusers • Compatible with the OpenAI client
|
|
47
|
+
</p>
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
<h1 align="center">Aquiles-Image</h1>
|
|
2
|
+
|
|
3
|
+
<div align="center">
|
|
4
|
+
<img src="aquilesimage/static/aquilesim.png" alt="Aquiles-Image Logo" width="225"/>
|
|
5
|
+
</div>
|
|
6
|
+
|
|
7
|
+
<p align="center">
|
|
8
|
+
<strong>Easy, fast and cheap Diffusion Models that work for everyone.</strong><br/>
|
|
9
|
+
🚀 FastAPI • Diffusers • Compatible with the OpenAI client
|
|
10
|
+
</p>
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: aquiles-image
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Aquiles-Image.
|
|
5
|
+
Author-email: Aquiles-ai / Fredy <riveraaai200678@gmail.com>
|
|
6
|
+
License: Apache License 2.0
|
|
7
|
+
Project-URL: Homepage, https://github.com/Aquiles-ai/Aquiles-Image
|
|
8
|
+
Project-URL: Issues, https://github.com/Aquiles-ai/Aquiles-Image/issues
|
|
9
|
+
Keywords: fastapi,ai
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
15
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
16
|
+
Classifier: Operating System :: OS Independent
|
|
17
|
+
Classifier: Environment :: Web Environment
|
|
18
|
+
Classifier: Topic :: Software Development :: Build Tools
|
|
19
|
+
Requires-Python: >=3.8
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
Requires-Dist: torch
|
|
22
|
+
Requires-Dist: torchvision
|
|
23
|
+
Requires-Dist: transformers
|
|
24
|
+
Requires-Dist: sentencepiece
|
|
25
|
+
Requires-Dist: fastapi
|
|
26
|
+
Requires-Dist: click>=8.0.0
|
|
27
|
+
Requires-Dist: uvicorn
|
|
28
|
+
Requires-Dist: ftfy
|
|
29
|
+
Requires-Dist: accelerate
|
|
30
|
+
Requires-Dist: xformers
|
|
31
|
+
Requires-Dist: protobuf
|
|
32
|
+
Provides-Extra: dev
|
|
33
|
+
Requires-Dist: pytest; extra == "dev"
|
|
34
|
+
Requires-Dist: black; extra == "dev"
|
|
35
|
+
Requires-Dist: isort; extra == "dev"
|
|
36
|
+
Requires-Dist: mypy; extra == "dev"
|
|
37
|
+
|
|
38
|
+
<h1 align="center">Aquiles-Image</h1>
|
|
39
|
+
|
|
40
|
+
<div align="center">
|
|
41
|
+
<img src="aquilesimage/static/aquilesim.png" alt="Aquiles-Image Logo" width="225"/>
|
|
42
|
+
</div>
|
|
43
|
+
|
|
44
|
+
<p align="center">
|
|
45
|
+
<strong>Easy, fast and cheap Diffusion Models that work for everyone.</strong><br/>
|
|
46
|
+
🚀 FastAPI • Diffusers • Compatible with the OpenAI client
|
|
47
|
+
</p>
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
aquiles_image.egg-info/PKG-INFO
|
|
4
|
+
aquiles_image.egg-info/SOURCES.txt
|
|
5
|
+
aquiles_image.egg-info/dependency_links.txt
|
|
6
|
+
aquiles_image.egg-info/entry_points.txt
|
|
7
|
+
aquiles_image.egg-info/requires.txt
|
|
8
|
+
aquiles_image.egg-info/top_level.txt
|
|
9
|
+
aquilesimage/main.py
|
|
10
|
+
aquilesimage/cli/__init__.py
|
|
11
|
+
aquilesimage/cli/cli.py
|
|
12
|
+
aquilesimage/configs/__init__.py
|
|
13
|
+
aquilesimage/configs/configs.py
|
|
14
|
+
aquilesimage/models/__init__.py
|
|
15
|
+
aquilesimage/models/models.py
|
|
16
|
+
aquilesimage/pipelines/__init__.py
|
|
17
|
+
aquilesimage/pipelines/pipelines.py
|
|
18
|
+
aquilesimage/runtime/__init__.py
|
|
19
|
+
aquilesimage/runtime/requestscopedpipeline.py
|
|
20
|
+
aquilesimage/runtime/scheduler.py
|
|
21
|
+
aquilesimage/static/aquilesim.png
|
|
22
|
+
aquilesimage/utils/__init__.py
|
|
23
|
+
aquilesimage/utils/utils.py
|
|
24
|
+
test/test.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
aquilesimage
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
from .cli import cli
|
|
@@ -0,0 +1,147 @@
|
|
|
1
|
+
import click
|
|
2
|
+
from typing import Optional
|
|
3
|
+
import sys
|
|
4
|
+
|
|
5
|
+
@click.group()
|
|
6
|
+
def cli():
|
|
7
|
+
"""A sample CLI application."""
|
|
8
|
+
pass
|
|
9
|
+
|
|
10
|
+
@cli.command("hello")
|
|
11
|
+
@click.option("--name")
|
|
12
|
+
def greet(name):
|
|
13
|
+
click.echo(f"Hello, {name}!")
|
|
14
|
+
|
|
15
|
+
@cli.command("serve")
|
|
16
|
+
@click.option("--host", default="0.0.0.0", help="Host where Aquiles-Image will be executed")
|
|
17
|
+
@click.option("--port", type=int, default=5500, help="Port where Aquiles-Image will be executed")
|
|
18
|
+
@click.option("--model", type=str, help="The model to use for image generation.")
|
|
19
|
+
@click.option("--api-key", type=str, help="API KEY enabled to make requests")
|
|
20
|
+
@click.option("--max-concurrent-infer", type=int, help="Maximum concurrent inferences")
|
|
21
|
+
@click.option("--block-request/--no-block-request", default=None, help="Block requests during maximum concurrent inferences")
|
|
22
|
+
@click.option("--force", is_flag=True, help="Force overwrite existing configuration")
|
|
23
|
+
def serve(host: str, port: int, model: Optional[str], api_key: Optional[str],
|
|
24
|
+
max_concurrent_infer: Optional[int], block_request: Optional[bool], force: bool):
|
|
25
|
+
"""Start the Aquiles-Image server."""
|
|
26
|
+
try:
|
|
27
|
+
import uvicorn
|
|
28
|
+
from aquilesimage.main import app
|
|
29
|
+
from aquilesimage.configs import load_config_cli, configs_image_serve
|
|
30
|
+
from aquilesimage.models import ConfigsServe
|
|
31
|
+
except ImportError as e:
|
|
32
|
+
click.echo(f"Error importing required modules: {e}", err=True)
|
|
33
|
+
sys.exit(1)
|
|
34
|
+
|
|
35
|
+
try:
|
|
36
|
+
conf = load_config_cli()
|
|
37
|
+
except Exception as e:
|
|
38
|
+
click.echo(f"Error loading configuration: {e}", err=True)
|
|
39
|
+
sys.exit(1)
|
|
40
|
+
|
|
41
|
+
model_from_config = conf.get("model")
|
|
42
|
+
final_model = model or model_from_config
|
|
43
|
+
|
|
44
|
+
if not final_model:
|
|
45
|
+
click.echo("Error: No model specified. Use --model parameter or configure one first.", err=True)
|
|
46
|
+
sys.exit(1)
|
|
47
|
+
|
|
48
|
+
config_needs_update = any([
|
|
49
|
+
model is not None,
|
|
50
|
+
api_key is not None,
|
|
51
|
+
max_concurrent_infer is not None,
|
|
52
|
+
block_request is not None
|
|
53
|
+
])
|
|
54
|
+
|
|
55
|
+
if config_needs_update:
|
|
56
|
+
try:
|
|
57
|
+
api_keys = [api_key] if api_key else []
|
|
58
|
+
|
|
59
|
+
gen_conf = ConfigsServe(
|
|
60
|
+
model=final_model,
|
|
61
|
+
allows_api_keys=api_keys,
|
|
62
|
+
max_concurrent_infer=max_concurrent_infer,
|
|
63
|
+
block_request=block_request
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
configs_image_serve(gen_conf, force)
|
|
67
|
+
click.echo("Configuration updated successfully.")
|
|
68
|
+
|
|
69
|
+
except Exception as e:
|
|
70
|
+
click.echo(f"Error saving configuration: {e}", err=True)
|
|
71
|
+
sys.exit(1)
|
|
72
|
+
|
|
73
|
+
click.echo(f"Starting server with:")
|
|
74
|
+
click.echo(f" Host: {host}")
|
|
75
|
+
click.echo(f" Port: {port}")
|
|
76
|
+
click.echo(f" Model: {final_model}")
|
|
77
|
+
|
|
78
|
+
try:
|
|
79
|
+
uvicorn.run(app, host=host, port=port)
|
|
80
|
+
except Exception as e:
|
|
81
|
+
click.echo(f"Error starting server: {e}", err=True)
|
|
82
|
+
sys.exit(1)
|
|
83
|
+
|
|
84
|
+
@cli.command("configs")
|
|
85
|
+
@click.option("--show", is_flag=True, help="Show current configuration")
|
|
86
|
+
@click.option("--reset", is_flag=True, help="Reset configuration to defaults")
|
|
87
|
+
def configs(show: bool, reset: bool):
|
|
88
|
+
"""Manage Aquiles-Image configuration."""
|
|
89
|
+
try:
|
|
90
|
+
from aquilesimage.configs import load_config_cli, clear_config_cache
|
|
91
|
+
import json
|
|
92
|
+
except ImportError as e:
|
|
93
|
+
click.echo(f"Error importing required modules: {e}", err=True)
|
|
94
|
+
sys.exit(1)
|
|
95
|
+
|
|
96
|
+
if reset:
|
|
97
|
+
if click.confirm("Are you sure you want to reset the configuration?"):
|
|
98
|
+
try:
|
|
99
|
+
clear_config_cache()
|
|
100
|
+
click.echo("Configuration reset successfully.")
|
|
101
|
+
except Exception as e:
|
|
102
|
+
click.echo(f"Error resetting configuration: {e}", err=True)
|
|
103
|
+
return
|
|
104
|
+
|
|
105
|
+
if show:
|
|
106
|
+
try:
|
|
107
|
+
conf = load_config_cli()
|
|
108
|
+
if conf:
|
|
109
|
+
click.echo("Current configuration:")
|
|
110
|
+
click.echo(json.dumps(conf, indent=2, ensure_ascii=False))
|
|
111
|
+
else:
|
|
112
|
+
click.echo("No configuration found.")
|
|
113
|
+
except Exception as e:
|
|
114
|
+
click.echo(f"Error loading configuration: {e}", err=True)
|
|
115
|
+
return
|
|
116
|
+
|
|
117
|
+
ctx = click.get_current_context()
|
|
118
|
+
click.echo(ctx.get_help())
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
@cli.command("validate")
|
|
122
|
+
def validate():
|
|
123
|
+
"""Validate current configuration."""
|
|
124
|
+
try:
|
|
125
|
+
from aquilesimage.configs import load_config_cli
|
|
126
|
+
from aquilesimage.models import ConfigsServe
|
|
127
|
+
except ImportError as e:
|
|
128
|
+
click.echo(f"Error importing required modules: {e}", err=True)
|
|
129
|
+
sys.exit(1)
|
|
130
|
+
|
|
131
|
+
try:
|
|
132
|
+
conf = load_config_cli()
|
|
133
|
+
|
|
134
|
+
if not conf:
|
|
135
|
+
click.echo("❌ No configuration found.", err=True)
|
|
136
|
+
sys.exit(1)
|
|
137
|
+
|
|
138
|
+
validated_conf = ConfigsServe(**conf)
|
|
139
|
+
click.echo("✅ Configuration is valid.")
|
|
140
|
+
|
|
141
|
+
except Exception as e:
|
|
142
|
+
click.echo(f"❌ Configuration validation failed: {e}", err=True)
|
|
143
|
+
sys.exit(1)
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
if __name__ == "__main__":
|
|
147
|
+
cli()
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
from .configs import configs_image_serve, load_config_cli, load_config_app, clear_config_cache
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
from platformdirs import user_data_dir
|
|
2
|
+
import json
|
|
3
|
+
import aiofiles
|
|
4
|
+
import asyncio
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
import os
|
|
7
|
+
from aquilesimage.models import ConfigsServe
|
|
8
|
+
from typing import Dict, Any
|
|
9
|
+
import time
|
|
10
|
+
import threading
|
|
11
|
+
from typing import Optional
|
|
12
|
+
|
|
13
|
+
_load_lock = asyncio.Lock()
|
|
14
|
+
data_dir = user_data_dir("aquiles", "Aquiles-Image")
|
|
15
|
+
os.makedirs(data_dir, exist_ok=True)
|
|
16
|
+
AQUILES_CONFIG = os.path.join(data_dir, "aquiles_cofig.json")
|
|
17
|
+
_cache_lock = threading.Lock()
|
|
18
|
+
_cached_config: Optional[Dict[str, Any]] = None
|
|
19
|
+
_cache_timestamp: float = 0
|
|
20
|
+
_cache_mtime: float = 0
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def load_config_cli(use_cache: bool = True, cache_ttl: float = 30.0) -> Dict[str, Any]:
|
|
24
|
+
global _cached_config, _cache_timestamp, _cache_mtime
|
|
25
|
+
config_path = Path(AQUILES_CONFIG)
|
|
26
|
+
if not config_path.exists():
|
|
27
|
+
return {}
|
|
28
|
+
current_time = time.time()
|
|
29
|
+
if use_cache:
|
|
30
|
+
with _cache_lock:
|
|
31
|
+
try:
|
|
32
|
+
file_mtime = config_path.stat().st_mtime
|
|
33
|
+
|
|
34
|
+
if (_cached_config is not None and
|
|
35
|
+
(current_time - _cache_timestamp) < cache_ttl and
|
|
36
|
+
file_mtime == _cache_mtime):
|
|
37
|
+
return _cached_config.copy()
|
|
38
|
+
|
|
39
|
+
except OSError:
|
|
40
|
+
pass
|
|
41
|
+
try:
|
|
42
|
+
with open(config_path, "r", encoding="utf-8") as f:
|
|
43
|
+
config_data = json.load(f)
|
|
44
|
+
|
|
45
|
+
if use_cache:
|
|
46
|
+
with _cache_lock:
|
|
47
|
+
try:
|
|
48
|
+
file_mtime = config_path.stat().st_mtime
|
|
49
|
+
_cached_config = config_data.copy()
|
|
50
|
+
_cache_timestamp = current_time
|
|
51
|
+
_cache_mtime = file_mtime
|
|
52
|
+
except OSError:
|
|
53
|
+
pass
|
|
54
|
+
|
|
55
|
+
return config_data
|
|
56
|
+
|
|
57
|
+
except FileNotFoundError:
|
|
58
|
+
return {}
|
|
59
|
+
except (json.JSONDecodeError, OSError, UnicodeDecodeError):
|
|
60
|
+
return {}
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
async def load_config_app() -> Dict[str, Any]:
|
|
64
|
+
async with _load_lock:
|
|
65
|
+
try:
|
|
66
|
+
async with aiofiles.open(AQUILES_CONFIG, "r", encoding="utf-8") as f:
|
|
67
|
+
s = await f.read()
|
|
68
|
+
except FileNotFoundError:
|
|
69
|
+
return {}
|
|
70
|
+
except Exception as exc:
|
|
71
|
+
return {}
|
|
72
|
+
|
|
73
|
+
try:
|
|
74
|
+
return json.loads(s)
|
|
75
|
+
except json.JSONDecodeError:
|
|
76
|
+
return {}
|
|
77
|
+
|
|
78
|
+
def clear_config_cache() -> None:
|
|
79
|
+
global _cached_config, _cache_timestamp, _cache_mtime
|
|
80
|
+
|
|
81
|
+
with _cache_lock:
|
|
82
|
+
_cached_config = None
|
|
83
|
+
_cache_timestamp = 0
|
|
84
|
+
_cache_mtime = 0
|
|
85
|
+
|
|
86
|
+
def configs_image_serve(cfg: ConfigsServe, force: bool = False) -> None:
|
|
87
|
+
conf = cfg.model_dump()
|
|
88
|
+
config_path = Path(AQUILES_CONFIG)
|
|
89
|
+
config_path.parent.mkdir(parents=True, exist_ok=True)
|
|
90
|
+
|
|
91
|
+
if config_path.exists() and not force:
|
|
92
|
+
return
|
|
93
|
+
try:
|
|
94
|
+
with open(config_path, "w", encoding="utf-8") as f:
|
|
95
|
+
json.dump(conf, f, ensure_ascii=False, indent=2)
|
|
96
|
+
|
|
97
|
+
clear_config_cache()
|
|
98
|
+
|
|
99
|
+
except (OSError, UnicodeEncodeError) as e:
|
|
100
|
+
pass
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
"""
|
|
2
|
+
The goal is to create image generation, editing, and variance endpoints compatible with the OpenAI client.
|
|
3
|
+
|
|
4
|
+
APIs:
|
|
5
|
+
|
|
6
|
+
POST /images/variations (create_variation)
|
|
7
|
+
POST /images/edits (edit)
|
|
8
|
+
POST /images/generations (generate)
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from fastapi import FastAPI, UploadFile, File, Request
|
|
12
|
+
from fastapi.middleware.cors import CORSMiddleware
|
|
13
|
+
from fastapi.concurrency import run_in_threadpool
|
|
14
|
+
from aquilesimage.models import CreateImageRequest, ImagesResponse, CreateImageEditRequest, CreateImageVariationRequest
|
|
15
|
+
from aquilesimage.utils import Utils, setup_colored_logger
|
|
16
|
+
from aquilesimage.runtime import RequestScopedPipeline
|
|
17
|
+
from aquilesimage.pipelines import ModelPipelineInit
|
|
18
|
+
from aquilesimage.configs import load_config_app, load_config_cli
|
|
19
|
+
import asyncio
|
|
20
|
+
import logging
|
|
21
|
+
from contextlib import asynccontextmanager
|
|
22
|
+
import threading
|
|
23
|
+
import torch
|
|
24
|
+
import random
|
|
25
|
+
|
|
26
|
+
logger = setup_colored_logger("Aquiles-Image", logging.INFO)
|
|
27
|
+
|
|
28
|
+
logger.info("Loading the model...")
|
|
29
|
+
|
|
30
|
+
model_pipeline = None
|
|
31
|
+
request_pipe = None
|
|
32
|
+
pipeline_lock = threading.Lock()
|
|
33
|
+
initializer = None
|
|
34
|
+
config = None
|
|
35
|
+
|
|
36
|
+
def load_models():
|
|
37
|
+
global model_pipeline, request_pipe, initializer, config
|
|
38
|
+
|
|
39
|
+
logger.info("Loading configuration...")
|
|
40
|
+
|
|
41
|
+
config = load_config_cli()
|
|
42
|
+
model_name = config.get("model")
|
|
43
|
+
|
|
44
|
+
if not model_name:
|
|
45
|
+
raise ValueError("No model specified in configuration. Please configure a model first.")
|
|
46
|
+
|
|
47
|
+
logger.info(f"Loading model: {model_name}")
|
|
48
|
+
|
|
49
|
+
try:
|
|
50
|
+
initializer = ModelPipelineInit(model=model_name)
|
|
51
|
+
model_pipeline = initializer.initialize_pipeline()
|
|
52
|
+
model_pipeline.start()
|
|
53
|
+
|
|
54
|
+
request_pipe = RequestScopedPipeline(model_pipeline.pipeline)
|
|
55
|
+
|
|
56
|
+
logger.info(f"Model '{model_name}' loaded successfully")
|
|
57
|
+
|
|
58
|
+
except Exception as e:
|
|
59
|
+
logger.error(f"Failed to initialize model pipeline: {e}")
|
|
60
|
+
raise
|
|
61
|
+
|
|
62
|
+
try:
|
|
63
|
+
load_models()
|
|
64
|
+
except Exception as e:
|
|
65
|
+
logger.error(f"Failed to initialize models: {e}")
|
|
66
|
+
raise
|
|
67
|
+
|
|
68
|
+
@asynccontextmanager
|
|
69
|
+
async def lifespan(app: FastAPI):
|
|
70
|
+
app.state.total_requests = 0
|
|
71
|
+
app.state.active_inferences = 0
|
|
72
|
+
app.state.metrics_lock = asyncio.Lock()
|
|
73
|
+
app.state.metrics_task = None
|
|
74
|
+
app.state.config = await load_config_app()
|
|
75
|
+
|
|
76
|
+
app.state.MODEL_INITIALIZER = initializer
|
|
77
|
+
app.state.MODEL_PIPELINE = model_pipeline
|
|
78
|
+
app.state.REQUEST_PIPE = request_pipe
|
|
79
|
+
app.state.PIPELINE_LOCK = pipeline_lock
|
|
80
|
+
|
|
81
|
+
# dumb config
|
|
82
|
+
app.state.utils_app = Utils(
|
|
83
|
+
host="0.0.0.0",
|
|
84
|
+
port=5500,
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
async def metrics_loop():
|
|
88
|
+
try:
|
|
89
|
+
while True:
|
|
90
|
+
async with app.state.metrics_lock:
|
|
91
|
+
total = app.state.total_requests
|
|
92
|
+
active = app.state.active_inferences
|
|
93
|
+
logger.info(f"[METRICS] total_requests={total} active_inferences={active}")
|
|
94
|
+
await asyncio.sleep(5)
|
|
95
|
+
except asyncio.CancelledError:
|
|
96
|
+
logger.info("Metrics loop cancelled")
|
|
97
|
+
raise
|
|
98
|
+
|
|
99
|
+
app.state.metrics_task = asyncio.create_task(metrics_loop())
|
|
100
|
+
|
|
101
|
+
try:
|
|
102
|
+
yield
|
|
103
|
+
finally:
|
|
104
|
+
task = app.state.metrics_task
|
|
105
|
+
if task:
|
|
106
|
+
task.cancel()
|
|
107
|
+
try:
|
|
108
|
+
await task
|
|
109
|
+
except asyncio.CancelledError:
|
|
110
|
+
pass
|
|
111
|
+
|
|
112
|
+
try:
|
|
113
|
+
stop_fn = getattr(model_pipeline, "stop", None) or getattr(model_pipeline, "close", None)
|
|
114
|
+
if callable(stop_fn):
|
|
115
|
+
await run_in_threadpool(stop_fn)
|
|
116
|
+
except Exception as e:
|
|
117
|
+
logger.warning(f"Error during pipeline shutdown: {e}")
|
|
118
|
+
|
|
119
|
+
if model_pipeline:
|
|
120
|
+
try:
|
|
121
|
+
stop_fn = getattr(model_pipeline, "stop", None) or getattr(model_pipeline, "close", None)
|
|
122
|
+
if callable(stop_fn):
|
|
123
|
+
await run_in_threadpool(stop_fn)
|
|
124
|
+
logger.info("Model pipeline stopped successfully")
|
|
125
|
+
except Exception as e:
|
|
126
|
+
logger.warning(f"Error during pipeline shutdown: {e}")
|
|
127
|
+
|
|
128
|
+
logger.info("Lifespan shutdown complete")
|
|
129
|
+
|
|
130
|
+
app = FastAPI(title="Aquiles-Image", lifespan=lifespan)
|
|
131
|
+
|
|
132
|
+
@app.middleware("http")
|
|
133
|
+
async def count_requests_middleware(request: Request, call_next):
|
|
134
|
+
async with app.state.metrics_lock:
|
|
135
|
+
app.state.total_requests += 1
|
|
136
|
+
response = await call_next(request)
|
|
137
|
+
return response
|
|
138
|
+
|
|
139
|
+
@app.post("/images/generations", response_model=ImagesResponse, tags=["Generation"])
|
|
140
|
+
async def create_image(input_r: CreateImageRequest):
|
|
141
|
+
def make_generator():
|
|
142
|
+
g = torch.Generator(device=initializer.device)
|
|
143
|
+
return g.manual_seed(random.randint(0, 10_000_000))
|
|
144
|
+
|
|
145
|
+
req_pipe = app.state.REQUEST_PIPE
|
|
146
|
+
pass
|
|
147
|
+
|
|
148
|
+
@app.post("/images/edits", response_model=ImagesResponse, tags=["Edit"])
|
|
149
|
+
async def create_image_edit(input_r: CreateImageEditRequest, image: UploadFile = File(...), mask: UploadFile | None = File(default=None)):
|
|
150
|
+
def make_generator():
|
|
151
|
+
g = torch.Generator(device=initializer.device)
|
|
152
|
+
return g.manual_seed(random.randint(0, 10_000_000))
|
|
153
|
+
|
|
154
|
+
req_pipe = app.state.REQUEST_PIPE
|
|
155
|
+
pass
|
|
156
|
+
|
|
157
|
+
@app.post("/images/variations", response_model=ImagesResponse, tags=["Variations"])
|
|
158
|
+
async def create_image_variation(input_r: CreateImageVariationRequest, image: UploadFile = File(...)):
|
|
159
|
+
def make_generator():
|
|
160
|
+
g = torch.Generator(device=initializer.device)
|
|
161
|
+
return g.manual_seed(random.randint(0, 10_000_000))
|
|
162
|
+
|
|
163
|
+
req_pipe = app.state.REQUEST_PIPE
|
|
164
|
+
pass
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
app.add_middleware(
|
|
168
|
+
CORSMiddleware,
|
|
169
|
+
allow_origins=["*"],
|
|
170
|
+
allow_credentials=True,
|
|
171
|
+
allow_methods=["*"],
|
|
172
|
+
allow_headers=["*"],
|
|
173
|
+
)
|
|
174
|
+
|
|
175
|
+
if __name__ == "__main__":
|
|
176
|
+
import uvicorn
|
|
177
|
+
|
|
178
|
+
uvicorn.run(app, host="0.0.0.0", port=5500)
|
|
@@ -0,0 +1,5 @@
|
|
|
1
|
+
from .models import (ImageModel, ResponseFormat, OutputFormat, BackgroundType, QualityType, StyleType,
|
|
2
|
+
ImageSize, InputFidelity, ModerationLevel, ImageGenInputUsageDetails, ImageGenUsage,
|
|
3
|
+
Image, ImagesResponse, ImageGenPartialImageEvent, ImagesUsage, ImageGenCompletedEvent,
|
|
4
|
+
ImageEditPartialImageEvent, ImageEditCompletedEvent, ImageGenStreamEvent, ImageEditStreamEvent,
|
|
5
|
+
CreateImageRequest, CreateImageEditRequest, CreateImageVariationRequest, ConfigsServe)
|