Vibe Transfer

playground/generate_image.py

@@ -11,7 +11,7 @@
 
 from novelai_python import APIError, Login
 from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
-from novelai_python.sdk.ai.generate_image import Action
+from novelai_python.sdk.ai.generate_image import Action, Sampler
 from novelai_python.utils.useful import enum_to_list
 
 load_dotenv()
@@ -30,6 +30,7 @@ async def main():
         gen = GenerateImageInfer.build(
             prompt=f"1girl, winter, jacket, sfw, angel, flower,{enhance}",
             action=Action.GENERATE,
+            sampler=Sampler.DDIM_V3
         )
         cost = gen.calculate_cost(is_opus=True)
         print(f"charge: {cost} if you are vip3")

playground/viber-transfer.py

@@ -0,0 +1,56 @@
+import asyncio
+import base64
+import os
+
+from dotenv import load_dotenv
+from pydantic import SecretStr
+
+from novelai_python import APIError, Login
+from novelai_python import GenerateImageInfer, ImageGenerateResp, JwtCredential
+from novelai_python.sdk.ai.generate_image import Action
+
+load_dotenv()
+
+enhance = "year 2023,dynamic angle, best quality, amazing quality, very aesthetic, absurdres"
+token = None
+jwt = os.getenv("NOVELAI_JWT") or token
+
+
+async def main():
+    jwt_class = JwtCredential(jwt_token=SecretStr(jwt))
+    _res = await Login.build(user_name=os.getenv("NOVELAI_USER"), password=os.getenv("NOVELAI_PASS")
+                             ).request()
+    with open("raw_test_image.png", "rb") as f:
+        data = f.read()
+    # Base64 encode the data
+    encoded = base64.b64encode(data).decode()
+    try:
+        gen = GenerateImageInfer.build(
+            prompt=f"1girl, spring, jacket, sfw, angel, flower,{enhance}",
+            action=Action.GENERATE,
+            image=encoded,
+            add_original_image=True,
+            strength=0.6,
+            reference_mode=True,  # IMPORTANT
+            width=1088,
+            height=896
+        )
+        cost = gen.calculate_cost(is_opus=True)
+        print(f"charge: {cost} if you are vip3")
+        print(f"charge: {gen.calculate_cost(is_opus=True)}")
+        _res = await gen.request(
+            session=jwt_class, remove_sign=True
+        )
+    except APIError as e:
+        print(e.response)
+        return
+
+    _res: ImageGenerateResp
+    print(_res.meta)
+    file = _res.files[0]
+    with open("generate_image_transfer.png", "wb") as f:
+        f.write(file[1])
+
+
+loop = asyncio.get_event_loop()
+loop.run_until_complete(main())

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.3.4"
+version = "0.3.5"
 description = "Novelai Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "[email protected]" },

src/novelai_python/sdk/ai/generate_image/__init__.py

@@ -34,7 +34,7 @@ class GenerateImageInfer(ApiBaseModel):
 
     class Params(BaseModel):
         # Inpaint
-        add_original_image: Optional[bool] = False
+        add_original_image: Optional[bool] = False  # FIXME: 未知作用
         mask: Optional[Union[str, bytes]] = None  # img2img,base64
 
         cfg_rescale: Optional[float] = Field(0, ge=0, le=1, multiple_of=0.02)
@@ -57,6 +57,20 @@ class Params(BaseModel):
 
         # Misc
         params_version: Optional[int] = 1
+        reference_image: Optional[Union[str, bytes]] = None
+        reference_information_extracted: Optional[float] = Field(1,
+                                                                 ge=0,
+                                                                 le=1,
+                                                                 multiple_of=0.01,
+                                                                 description="extracting concepts or features"
+                                                                 )
+        reference_strength: Optional[float] = Field(0.6,
+                                                    ge=0,
+                                                    le=1,
+                                                    multiple_of=0.01,
+                                                    description="the stronger the AI will try to emulate visual cues."
+                                                    )
+
         legacy: Optional[bool] = False
         legacy_v3_extend: Optional[bool] = False
 
@@ -84,14 +98,6 @@ class Params(BaseModel):
         """Undesired Content Strength"""
         width: Optional[int] = Field(832, ge=64, le=49152)
 
-        @model_validator(mode="after")
-        def validate_img2img(self):
-            image = True if self.image else False
-            add_origin = True if self.add_original_image else False
-            if image != add_origin:
-                raise ValueError('Invalid Model Params For img2img2 mode... image should match add_original_image!')
-            return self
-
         @field_validator('mask')
         def mask_validator(cls, v: Union[str, bytes]):
             if isinstance(v, str) and v.startswith("data:image/"):
@@ -108,6 +114,14 @@ def image_validator(cls, v: Union[str, bytes]):
                 return base64.b64encode(v).decode("utf-8")
             return v
 
+        @field_validator('reference_image')
+        def image_validator(cls, v: Union[str, bytes]):
+            if isinstance(v, str) and v.startswith("data:image/"):
+                raise ValueError("Invalid image format, must be base64 encoded.")
+            if isinstance(v, bytes):
+                return base64.b64encode(v).decode("utf-8")
+            return v
+
         @field_validator('width')
         def width_validator(cls, v: int):
             """
@@ -176,13 +190,20 @@ def model_post_init(self, *args) -> None:
 
     @model_validator(mode="after")
     def validate_model(self):
-        if self.action == Action.INFILL and not self.parameters.mask:
-            logger.warning("Mask maybe required for infill mode.")
+        if self.action == Action.INFILL:
+            if not self.parameters.mask:
+                logger.warning("Mask maybe required for infill mode.")
         if self.action != Action.GENERATE:
             self.parameters.extra_noise_seed = self.parameters.seed
         if self.action == Action.IMG2IMG:
             self.parameters.sm = False
             self.parameters.sm_dyn = False
+            if not self.parameters.image:
+                raise ValueError("image is must required for img2img mode.")
+            if not self.parameters.add_original_image:
+                raise ValueError("add_original_image is must required for img2img mode.")
+        if self.parameters.image and self.parameters.reference_image:
+            logger.warning("image and reference_image should not be used together.")
         return self
 
     @property
@@ -247,6 +268,7 @@ def build(cls,
               qualityToggle: bool = True,
               ucPreset: Union[UCPreset, int] = UCPreset.TYPE0,
               image: Union[str, bytes] = None,
+              reference_mode: bool = False,
               add_original_image: bool = None,
               strength: float = None,
               mask: Union[str, bytes] = None,
@@ -264,12 +286,13 @@ def build(cls,
         :param negative_prompt: The content of negative prompt
         :param seed: The seed for generate image
         :param steps: The steps for generate image
-        :param scale: Prompt Guidance
+        :param scale: The scale for generate image
         :param cfg_rescale: Prompt Guidance Rescale 0-1 lower is more creative
         :param sampler: The sampler for generate image
         :param width: 宽
         :param height: 高
         :param image: 图片
+        :param reference_mode: 是否是参考模式
         :param add_original_image: 是否添加原始图片
         :param strength: IMG2IMG 强度
         :param mask: Inpainting mask
@@ -290,13 +313,21 @@ def build(cls,
             "height": height,
             "qualityToggle": qualityToggle,
             "ucPreset": ucPreset,
-            "image": image,
             "add_original_image": add_original_image,
-            "strength": strength,
             "mask": mask,
             "controlnet_model": controlnet_model,
             "controlnet_condition": controlnet_condition
         })
+        if reference_mode:
+            kwargs.update({
+                "reference_image": image,
+                "reference_strength": strength,
+            })
+        else:
+            kwargs.update({
+                "image": image,
+                "strength": strength,
+            })
         # 清理空值
         param = {k: v for k, v in kwargs.items() if v is not None}
         return cls(
@@ -361,17 +392,15 @@ async def request(self,
         try:
             _log_data = deepcopy(request_data)
             if self.action == Action.GENERATE:
+                _log_data.get("parameters", {}).update({
+                    "reference_image": "base64 data" if self.parameters.reference_image else None,
+                })
                 logger.debug(f"Request Data: {_log_data}")
             else:
                 _log_data.get("parameters", {}).update({
                     "image": "base64 data" if self.parameters.image else None,
-                }
-                )
-                _log_data.get("parameters", {}).update(
-                    {
-                        "mask": "base64 data" if self.parameters.mask else None,
-                    }
-                )
+                    "mask": "base64 data" if self.parameters.mask else None,
+                })
                 logger.debug(f"Request Data: {_log_data}")
             del _log_data
         except Exception as e:
@@ -383,7 +412,7 @@ async def request(self,
                 data=json.dumps(request_data).encode("utf-8")
             )
             if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

src/novelai_python/sdk/ai/generate_image/suggest_tags.py

@@ -84,7 +84,7 @@ async def request(self,
                 url=self.base_url + "?" + "&".join([f"{k}={v}" for k, v in request_data.items()])
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

src/novelai_python/sdk/ai/upscale.py

@@ -131,7 +131,7 @@ async def request(self,
                 data=json.dumps(request_data).encode("utf-8")
             )
             if response.headers.get('Content-Type') not in ['binary/octet-stream', 'application/x-zip-compressed']:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

src/novelai_python/sdk/user/information.py

@@ -73,7 +73,7 @@ async def request(self,
                 self.base_url
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

src/novelai_python/sdk/user/login.py

@@ -92,7 +92,7 @@ async def request(self,
                 data=json.dumps(request_data).encode("utf-8")
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 201:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try:

src/novelai_python/sdk/user/subscription.py

@@ -80,7 +80,7 @@ async def request(self,
                 url=self.base_url,
             )
             if "application/json" not in response.headers.get('Content-Type') or response.status_code != 200:
-                logger.error(
+                logger.warning(
                     f"Error with content type: {response.headers.get('Content-Type')} and code: {response.status_code}"
                 )
                 try: