Merge pull request #3722 from evshiron/feat/progress-api

prototype progress api

1 files changed, 20 insertions, 11 deletions

diff --git a/modules/api/models.py b/modules/api/models.py
index 58e8e58b..9ee42a17 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -52,17 +52,17 @@ class PydanticModelGenerator:
             # field_type = str if not overrides.get(k) else overrides[k]["type"]
             # print(k, v.annotation, v.default)
             field_type = v.annotation
-            
+
             return Optional[field_type]
-        
+
         def merge_class_params(class_):
             all_classes = list(filter(lambda x: x is not object, inspect.getmro(class_)))
             parameters = {}
             for classes in all_classes:
                 parameters = {**parameters, **inspect.signature(classes.__init__).parameters}
             return parameters
-            
-                
+
+
         self._model_name = model_name
         self._class_data = merge_class_params(class_instance)
         self._model_def = [
@@ -74,11 +74,11 @@ class PydanticModelGenerator:
             )
             for (k,v) in self._class_data.items() if k not in API_NOT_ALLOWED
         ]
-        
+
         for fields in additional_fields:
             self._model_def.append(ModelDef(
-                field=underscore(fields["key"]), 
-                field_alias=fields["key"], 
+                field=underscore(fields["key"]),
+                field_alias=fields["key"],
                 field_type=fields["type"],
                 field_value=fields["default"],
                 field_exclude=fields["exclude"] if "exclude" in fields else False))
@@ -95,15 +95,15 @@ class PydanticModelGenerator:
         DynamicModel.__config__.allow_population_by_field_name = True
         DynamicModel.__config__.allow_mutation = True
         return DynamicModel
-    
+
 StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
-    "StableDiffusionProcessingTxt2Img", 
+    "StableDiffusionProcessingTxt2Img",
     StableDiffusionProcessingTxt2Img,
     [{"key": "sampler_index", "type": str, "default": "Euler"}]
 ).generate_model()
 
 StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
-    "StableDiffusionProcessingImg2Img", 
+    "StableDiffusionProcessingImg2Img",
     StableDiffusionProcessingImg2Img,
     [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "init_images", "type": list, "default": None}, {"key": "denoising_strength", "type": float, "default": 0.75}, {"key": "mask", "type": str, "default": None}, {"key": "include_init_images", "type": bool, "default": False, "exclude" : True}]
 ).generate_model()
@@ -155,4 +155,13 @@ class PNGInfoRequest(BaseModel):
     image: str = Field(title="Image", description="The base64 encoded PNG image")
 
 class PNGInfoResponse(BaseModel):
-    info: str = Field(title="Image info", description="A string with all the info the image had")
\ No newline at end of file
+    info: str = Field(title="Image info", description="A string with all the info the image had")
+
+class ProgressRequest(BaseModel):
+    skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization")
+
+class ProgressResponse(BaseModel):
+    progress: float = Field(title="Progress", description="The progress with a range of 0 to 1")
+    eta_relative: float = Field(title="ETA in secs")
+    state: dict = Field(title="State", description="The current state snapshot")
+    current_image: str = Field(default=None, title="Current image", description="The current image in base64 format. opts.show_progress_every_n_steps is required for this to work.")