Spaces:
Sleeping
Sleeping
updage model registry
Browse files
app.py
CHANGED
|
@@ -17,6 +17,8 @@ import functools
|
|
| 17 |
|
| 18 |
app = FastAPI()
|
| 19 |
|
|
|
|
|
|
|
| 20 |
# Comprehensive model registry
|
| 21 |
MODELS = {
|
| 22 |
"SDXL-Base": {
|
|
@@ -30,9 +32,53 @@ MODELS = {
|
|
| 30 |
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 31 |
}
|
| 32 |
},
|
| 33 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 34 |
}
|
| 35 |
|
|
|
|
| 36 |
class ModelManager:
|
| 37 |
def __init__(self):
|
| 38 |
self.current_model = None
|
|
@@ -60,7 +106,8 @@ class ModelManager:
|
|
| 60 |
model_id,
|
| 61 |
torch_dtype=self._dtype,
|
| 62 |
variant="fp16" if self._device == "cuda" else None,
|
| 63 |
-
device_map="
|
|
|
|
| 64 |
)
|
| 65 |
|
| 66 |
def load_model(self, model_name: str):
|
|
|
|
| 17 |
|
| 18 |
app = FastAPI()
|
| 19 |
|
| 20 |
+
# Comprehensive model registry
|
| 21 |
+
|
| 22 |
# Comprehensive model registry
|
| 23 |
MODELS = {
|
| 24 |
"SDXL-Base": {
|
|
|
|
| 32 |
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 33 |
}
|
| 34 |
},
|
| 35 |
+
"SDXL-Turbo": {
|
| 36 |
+
"model_id": "stabilityai/sdxl-turbo",
|
| 37 |
+
"pipeline": AutoPipelineForText2Image,
|
| 38 |
+
"supports_img2img": True,
|
| 39 |
+
"parameters": {
|
| 40 |
+
"num_inference_steps": {"min": 1, "max": 50, "default": 1},
|
| 41 |
+
"guidance_scale": {"min": 0.0, "max": 20.0, "default": 7.5},
|
| 42 |
+
"width": {"min": 256, "max": 1024, "default": 512, "step": 64},
|
| 43 |
+
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 44 |
+
}
|
| 45 |
+
},
|
| 46 |
+
"SD-1.5": {
|
| 47 |
+
"model_id": "runwayml/stable-diffusion-v1-5",
|
| 48 |
+
"pipeline": StableDiffusionPipeline,
|
| 49 |
+
"supports_img2img": True,
|
| 50 |
+
"parameters": {
|
| 51 |
+
"num_inference_steps": {"min": 1, "max": 50, "default": 30},
|
| 52 |
+
"guidance_scale": {"min": 1, "max": 20, "default": 7.5},
|
| 53 |
+
"width": {"min": 256, "max": 1024, "default": 512, "step": 64},
|
| 54 |
+
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 55 |
+
}
|
| 56 |
+
},
|
| 57 |
+
"Waifu-Diffusion": {
|
| 58 |
+
"model_id": "hakurei/waifu-diffusion",
|
| 59 |
+
"pipeline": StableDiffusionPipeline,
|
| 60 |
+
"supports_img2img": True,
|
| 61 |
+
"parameters": {
|
| 62 |
+
"num_inference_steps": {"min": 1, "max": 100, "default": 50},
|
| 63 |
+
"guidance_scale": {"min": 1, "max": 15, "default": 7.5},
|
| 64 |
+
"width": {"min": 256, "max": 1024, "default": 512, "step": 64},
|
| 65 |
+
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 66 |
+
}
|
| 67 |
+
},
|
| 68 |
+
"Flux": {
|
| 69 |
+
"model_id": "black-forest-labs/flux-1-1-dev",
|
| 70 |
+
"pipeline": AutoPipelineForText2Image,
|
| 71 |
+
"supports_img2img": True,
|
| 72 |
+
"parameters": {
|
| 73 |
+
"num_inference_steps": {"min": 1, "max": 50, "default": 25},
|
| 74 |
+
"guidance_scale": {"min": 1, "max": 15, "default": 7.5},
|
| 75 |
+
"width": {"min": 256, "max": 1024, "default": 512, "step": 64},
|
| 76 |
+
"height": {"min": 256, "max": 1024, "default": 512, "step": 64}
|
| 77 |
+
}
|
| 78 |
+
}
|
| 79 |
}
|
| 80 |
|
| 81 |
+
|
| 82 |
class ModelManager:
|
| 83 |
def __init__(self):
|
| 84 |
self.current_model = None
|
|
|
|
| 106 |
model_id,
|
| 107 |
torch_dtype=self._dtype,
|
| 108 |
variant="fp16" if self._device == "cuda" else None,
|
| 109 |
+
device_map="balanced"
|
| 110 |
+
|
| 111 |
)
|
| 112 |
|
| 113 |
def load_model(self, model_name: str):
|