Merge pull request #11 from pokedim13/dev

Dev
pokedim13 · Oct 7, 2024 · 89e3139 · 89e3139
2 parents 56d74aa + 620f99b
commit 89e3139
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wombo"
-version = "0.3.0"
+version = "0.3.2"
 description = "AI picture generated tool"
 authors = ["Алексей Баев <baev@cumproject.ru>"]
 readme = "README.md"

diff --git a/tests/example.py b/tests/example.py
@@ -3,9 +3,7 @@
 
 async def main():
     dream = AsyncDream("eyJhbGciOiJSUzI1NiIsImtpZCI6IjhkNzU2OWQyODJkNWM1Mzk5MmNiYWZjZWI2NjBlYmQ0Y2E1OTMxM2EiLCJ0eXAiOiJKV1QifQ.eyJwcm92aWRlcl9pZCI6ImFub255bW91cyIsImlzcyI6Imh0dHBzOi8vc2VjdXJldG9rZW4uZ29vZ2xlLmNvbS9wYWludC1wcm9kIiwiYXVkIjoicGFpbnQtcHJvZCIsImF1dGhfdGltZSI6MTcyNzE4NzIyNSwidXNlcl9pZCI6IlRBeXZseU5FMmRSVlh3YmFDY1IyZGxZcmFhMDIiLCJzdWIiOiJUQXl2bHlORTJkUlZYd2JhQ2NSMmRsWXJhYTAyIiwiaWF0IjoxNzI4MTg3MTQwLCJleHAiOjE3MjgxOTA3NDAsImZpcmViYXNlIjp7ImlkZW50aXRpZXMiOnt9LCJzaWduX2luX3Byb3ZpZGVyIjoiYW5vbnltb3VzIn19.pFLDKAIU7XaXdfPTEJQJFrKK-FE799Pooiw-KRUk2y6lHgu_GibSatUwp_TNVg1T3sXSLbKCgtlVVEZUfTtxD26ckuapQ7UX7zOtLOUg_Ou_QvWtgceoJOoLDX1voYhnz5gEX65nCLGvR5t39gqiqIHsAhHjYMLMk77JvB7dqy4x4t6KCFkJCs8mCN45pnVqQCmdixZvY0Ers20gTIk_mCqaD7x1H9Pba4OL5Ro3qZEjpz8O_4Bwm1IY4YBqJepk5idxJ_emVEjfufDOVL8HOHNJBqsoLWz02unumUSVHgJoHr5MZyLWeKVelaLpB6LdT-sCoJ4xtmARJb07j2is-A", debug=True)
-    print(await dream.generate("anime waifu"))
+    print(await dream.style._get_styles())
+    print(len(dream.style.free.dict()))
 
-
-# dream = Dream("eyJhbGciOiJSUzI1NiIsImtpZCI6IjhkNzU2OWQyODJkNWM1Mzk5MmNiYWZjZWI2NjBlYmQ0Y2E1OTMxM2EiLCJ0eXAiOiJKV1QifQ.eyJwcm92aWRlcl9pZCI6ImFub255bW91cyIsImlzcyI6Imh0dHBzOi8vc2VjdXJldG9rZW4uZ29vZ2xlLmNvbS9wYWludC1wcm9kIiwiYXVkIjoicGFpbnQtcHJvZCIsImF1dGhfdGltZSI6MTcyNzE4NzIyNSwidXNlcl9pZCI6IlRBeXZseU5FMmRSVlh3YmFDY1IyZGxZcmFhMDIiLCJzdWIiOiJUQXl2bHlORTJkUlZYd2JhQ2NSMmRsWXJhYTAyIiwiaWF0IjoxNzI4MTg3MTQwLCJleHAiOjE3MjgxOTA3NDAsImZpcmViYXNlIjp7ImlkZW50aXRpZXMiOnt9LCJzaWduX2luX3Byb3ZpZGVyIjoiYW5vbnltb3VzIn19.pFLDKAIU7XaXdfPTEJQJFrKK-FE799Pooiw-KRUk2y6lHgu_GibSatUwp_TNVg1T3sXSLbKCgtlVVEZUfTtxD26ckuapQ7UX7zOtLOUg_Ou_QvWtgceoJOoLDX1voYhnz5gEX65nCLGvR5t39gqiqIHsAhHjYMLMk77JvB7dqy4x4t6KCFkJCs8mCN45pnVqQCmdixZvY0Ers20gTIk_mCqaD7x1H9Pba4OL5Ro3qZEjpz8O_4Bwm1IY4YBqJepk5idxJ_emVEjfufDOVL8HOHNJBqsoLWz02unumUSVHgJoHr5MZyLWeKVelaLpB6LdT-sCoJ4xtmARJb07j2is-A", debug=True)
-# print(dream.generate("anime waifu"))
 asyncio.run(main())
diff --git a/wombo/base/dream.py b/wombo/base/dream.py
@@ -21,6 +21,25 @@ def __getitem__(self, key):
                     return style.id 
             return None
 
+        @property
+        def free(self) -> StyleModel:
+            res = []
+            for style in self.styles.root:
+                if style.is_premium:
+                    continue
+                res.append(style)
+            return self.dream._get_model(StyleModel, res)
+
+        @property
+        def premium(self) -> StyleModel:
+            res = []
+            for style in self.styles.root:
+                if not style.is_premium:
+                    continue
+                res.append(style)
+            return self.dream._get_model(StyleModel, res)
+
+
         @abstractproperty
         def url(self) -> str:
             """Returns the URL associated with this styles.