Merge pull request #141 from VoltaML/experimental

v0.3.1 - bug fixes, non squashed merge for better track of changes

Author: Stax124

core/inference/pytorch/pipeline.py

@@ -719,6 +719,8 @@ def img2img(
         image: Union[torch.FloatTensor, PIL.Image.Image],  # type: ignore
         prompt: Union[str, List[str]],
         generator: torch.Generator,
+        height: int = 512,
+        width: int = 512,
         negative_prompt: Optional[Union[str, List[str]]] = None,
         strength: float = 0.8,
         num_inference_steps: Optional[int] = 50,
@@ -795,6 +797,8 @@ def img2img(
             prompt=prompt,
             negative_prompt=negative_prompt,
             image=image,
+            height=height,
+            width=width,
             num_inference_steps=num_inference_steps,  # type: ignore
             guidance_scale=guidance_scale,  # type: ignore
             self_attention_scale=self_attention_scale,

core/inference/pytorch/pytorch.py

@@ -310,6 +310,8 @@ def txt2img(self, job: Txt2ImgQueueEntry) -> List[Image.Image]:
                 data = pipe.img2img(
                     prompt=job.data.prompt,
                     image=latents,
+                    height=latents.shape[2] * 8,
+                    width=latents.shape[3] * 8,
                     num_inference_steps=flag.steps,
                     guidance_scale=job.data.guidance_scale,
                     self_attention_scale=job.data.self_attention_scale,
@@ -360,6 +362,8 @@ def img2img(self, job: Img2ImgQueueEntry) -> List[Image.Image]:
             data = pipe.img2img(
                 prompt=job.data.prompt,
                 image=input_image,
+                height=job.data.height,  # technically isn't needed, but it's here for consistency sake
+                width=job.data.width,  # technically isn't needed, but it's here for consistency sake
                 num_inference_steps=job.data.steps,
                 guidance_scale=job.data.guidance_scale,
                 self_attention_scale=job.data.self_attention_scale,

core/inference/utilities/latents.py

@@ -488,19 +488,13 @@ def scale_latents(
 ):
     "Interpolate the latents to the desired scale."
 
-    align_to = (
-        32 if latent_scale_mode in ["bislerp-tortured", "bislerp-original"] else 8
-    )
-
     s = time()
 
-    logger.debug(
-        f"Scaling latents with shape {list(latents.shape)}, scale: {scale}, alignment: {align_to}"
-    )
+    logger.debug(f"Scaling latents with shape {list(latents.shape)}, scale: {scale}")
 
     # Scale and round to multiple of 32
-    width_truncated = int(((latents.shape[2] * scale - 1) // align_to + 1) * align_to)
-    height_truncated = int(((latents.shape[3] * scale - 1) // align_to + 1) * align_to)
+    width_truncated = int(latents.shape[2] * scale)
+    height_truncated = int(latents.shape[3] * scale)
 
     # Scale the latents
     if latent_scale_mode == "bislerp-tortured":

core/shared.py

@@ -18,4 +18,4 @@
 
 uvicorn_server: Optional["Server"] = None
 uvicorn_loop: Optional[asyncio.AbstractEventLoop] = None
-asyncio_tasks: list[asyncio.Task] = []
+asyncio_tasks: List[asyncio.Task] = []