fix: pydantic models for http server working now. Fixes #380

.gitignore

@@ -30,3 +30,4 @@ tests/vastai_cli.py
 **/.eggs
 /img_size_memory_usage.csv
 /tests/test_cluster_output/
+/.env

README.md

@@ -493,6 +493,10 @@ A: The AI models are cached in `~/.cache/` (or `HUGGINGFACE_HUB_CACHE`). To dele
 
 ## ChangeLog
 
+**13.2.1**
+- fix: pydantic models for http server working now. Fixes #380
+- fix: install triton so annoying message is gone
+
 **13.2.0**
 - fix: allow tile_mode to be set to True or False for backward compatibility
 - fix: various pydantic issues have been resolved

imaginairy/http/stablestudio/models.py

@@ -119,5 +119,5 @@ class StableStudioBatchResponse(BaseModel):
     images: List[StableStudioImage]
 
 
-StableStudioInput.update_forward_refs()
-StableStudioImage.update_forward_refs()
+StableStudioInput.model_rebuild()
+StableStudioImage.model_rebuild()

imaginairy/http/utils.py

@@ -5,7 +5,7 @@
 
 
 def generate_image(prompt):
-    """ImaginPrompt to generated image."""
+    """ImaginePrompt to generated image."""
     result = next(imagine([prompt]))
     img = result.images["generated"]
     img_io = BytesIO()
@@ -27,7 +27,7 @@ def __get_validators__(cls):
         yield cls.validate
 
     @classmethod
-    def validate(cls, v):
+    def validate(cls, v, info):
         if isinstance(v, bytes):
             return v
         if isinstance(v, str):

imaginairy/schema.py

@@ -90,6 +90,12 @@ def __getattr__(self, key):
         self._load_img()
         return getattr(self._img, key)
 
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+    def __getstate__(self):
+        return self.__dict__
+
     def _load_img(self):
         if self._img is None:
             from PIL import Image, ImageOps
@@ -192,7 +198,10 @@ def __repr__(self):
 
         shows filepath or url if available.
         """
-        return f"<LazyLoadingImage filepath={self._lazy_filepath} url={self._lazy_url}>"
+        try:
+            return f"<LazyLoadingImage filepath={self._lazy_filepath} url={self._lazy_url}>"
+        except Exception as e:
+            return f"<LazyLoadingImage *{e}*>"
 
 
 #
@@ -245,7 +254,7 @@ def __repr__(self):
         return f"{self.weight}*({self.text})"
 
 
-class ImaginePrompt(BaseModel):
+class ImaginePrompt(BaseModel, protected_namespaces=()):
     prompt: Optional[List[WeightedPrompt]] = Field(default=None, validate_default=True)
     negative_prompt: Optional[List[WeightedPrompt]] = Field(
         default=None, validate_default=True

requirements-dev.txt

@@ -42,13 +42,13 @@ click-help-colors==0.9.2
     # via imaginAIry (setup.py)
 click-shell==2.1
     # via imaginAIry (setup.py)
-contourpy==1.1.0
+contourpy==1.1.1
     # via matplotlib
 coverage==7.3.1
     # via -r requirements-dev.in
 cycler==0.11.0
     # via matplotlib
-diffusers==0.20.2
+diffusers==0.21.1
     # via imaginAIry (setup.py)
 dill==0.3.7
     # via pylint
@@ -64,7 +64,7 @@ fairscale==0.4.13
     # via imaginAIry (setup.py)
 fastapi==0.103.1
     # via imaginAIry (setup.py)
-filelock==3.12.3
+filelock==3.12.4
     # via
     #   diffusers
     #   huggingface-hub
@@ -77,7 +77,7 @@ frozenlist==1.4.0
     # via
     #   aiohttp
     #   aiosignal
-fsspec[http]==2023.9.0
+fsspec[http]==2023.9.1
     # via
     #   huggingface-hub
     #   pytorch-lightning
@@ -122,7 +122,7 @@ lightning-utilities==0.9.0
     #   torchmetrics
 llvmlite==0.40.1
     # via numba
-matplotlib==3.7.3
+matplotlib==3.8.0
     # via filterpy
 mccabe==0.7.0
     # via
@@ -178,7 +178,7 @@ pathspec==0.11.2
     # via
     #   black
     #   pycln
-pillow==10.0.0
+pillow==10.0.1
     # via
     #   diffusers
     #   facexlib
@@ -257,7 +257,7 @@ requests==2.31.0
     #   transformers
 responses==0.23.3
     # via -r requirements-dev.in
-ruff==0.0.288
+ruff==0.0.290
     # via -r requirements-dev.in
 safetensors==0.3.3
     # via
@@ -330,7 +330,7 @@ tqdm==4.66.1
     #   open-clip-torch
     #   pytorch-lightning
     #   transformers
-transformers==4.33.1
+transformers==4.33.2
     # via imaginAIry (setup.py)
 typer==0.9.0
     # via pycln
@@ -341,7 +341,6 @@ typing-extensions==4.7.1
     #   astroid
     #   black
     #   fastapi
-    #   filelock
     #   huggingface-hub
     #   libcst
     #   lightning-utilities

setup.py

@@ -103,9 +103,10 @@ def get_git_revision_hash() -> str:
         "scipy<1.11",
         "timm>=0.4.12,!=0.9.0,!=0.9.1",  # for vendored blip
         "torchdiffeq>=0.2.0",
-        "transformers>=4.19.2",
         "torchmetrics>=0.6.0",
         "torchvision>=0.13.1",
+        "transformers>=4.19.2",
+        "triton>=2.0.0; sys_platform!='darwin' and platform_machine!='aarch64'",
         "kornia>=0.6",
         "uvicorn>=0.16.0",
         "xformers>=0.0.16; sys_platform!='darwin' and platform_machine!='aarch64'",

tests/test_schema/test_lazy_load_image.py

@@ -101,3 +101,14 @@ def test_image_deserialization(red_path, red_url):
     for row in rows:
         obj = TestModel.model_validate(row)
         assert obj.header_img.size == (512, 512)
+
+
+def test_image_state(red_path):
+    """I dont remember what this fixes. Maybe the ability of pydantic to copy an object?."""
+    img = LazyLoadingImage(filepath=red_path)
+
+    # bypass init
+    img2 = LazyLoadingImage.__new__(LazyLoadingImage)
+    img2.__setstate__(img.__getstate__())
+
+    assert repr(img) == repr(img2)