Image-to-image fixes

README.md

@@ -49,14 +49,13 @@ preload_from_hub:
 
 Gradio app for Stable Diffusion 1.5 including:
 * txt2img and img2img pipelines
-* curated models and TI embeddings
-* multiple samplers with Karras schedule
-* Compel prompting
+* Curated models and TI embeddings
 * 100+ styles from sdxl_prompt_styler
-* FreeU and Clip Skip for quality
-* DeepCache and ToMe for speed
+* Compel prompt weighting
+* Multiple samplers with Karras scheduling
+* DeepCache, ToMe, FreeU, and Clip Skip available
 * Real-ESRGAN upscaling
-* optional TAESD
+* Optional tiny autoencoder
 
 ## Usage
 
@@ -84,5 +83,4 @@ python cli.py 'an astronaut riding a horse on mars'
 
 ## TODO
 
-- [ ] Metadata embed and display
 - [ ] IP-Adapter and T2I-Adapter

app.py

@@ -48,11 +48,10 @@ def random_fn():
 def gallery_fn(images, image):
     if image is not None:
         return gr.Dropdown(
-            choices=[("🔒", -1)],
+            choices=[("🔒", -2)],
             interactive=False,
-            value=-1,
+            value=-2,
         )
-
     return gr.Dropdown(
         choices=[("None", -1)]
         + [(str(i + 1), i) for i, _ in enumerate(images if images is not None else [])],
@@ -67,9 +66,12 @@ def image_prompt_fn(images):
 
 # can't use image input in JS
 def image_select_fn(images, image, i):
-    if image is not None and i == -1:
-        return gr.Image(value=image)
-    return gr.Image(value=images[i][0]) if i > -1 else None
+    # -2 is the lock icon, -1 is None
+    if i == -2:
+        return gr.Image(image)
+    if i == -1:
+        return gr.Image(None)
+    return gr.Image(images[i][0]) if i > -1 else None
 
 
 def generate_fn(*args):

lib/inference.py

@@ -7,12 +7,14 @@ from datetime import datetime
 from itertools import product
 from typing import Callable
 
+import numpy as np
 import spaces
 import tomesd
 import torch
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
 from huggingface_hub.utils import HFValidationError, RepositoryNotFoundError
+from PIL import Image
 
 from .loader import Loader
 
@@ -66,6 +68,21 @@ def apply_style(prompt, style_id, negative=False):
     return prompt
 
 
+def prepare_image(input, size=(512, 512)):
+    image = None
+    if isinstance(input, Image.Image):
+        image = input
+    if isinstance(input, np.ndarray):
+        image = Image.fromarray(input)
+    if isinstance(input, str):
+        if os.path.isfile(input):
+            image = Image.open(input)
+    if image is not None:
+        return image.convert("RGB").resize(size, Image.Resampling.LANCZOS)
+    else:
+        raise ValueError("Invalid image prompt")
+
+
 @spaces.GPU(duration=40)
 def generate(
     positive_prompt,
@@ -196,8 +213,8 @@ def generate(
             }
 
             if KIND == "img2img":
-                kwargs["image"] = image_prompt
                 kwargs["strength"] = denoising_strength
+                kwargs["image"] = prepare_image(image_prompt, (width, height))
 
             with token_merging(pipe, tome_ratio=tome_ratio):
                 try:

lib/loader.py

@@ -72,7 +72,7 @@ class Loader:
             print("Switching to Tiny VAE...")
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 pretrained_model_name_or_path="madebyollin/taesd",
-            ).to(self.pipe.device)
+            ).to(self.pipe.device, self.pipe.dtype)
             return
 
         if is_tiny and not taesd:
@@ -81,7 +81,7 @@ class Loader:
                 pretrained_model_name_or_path=model_name,
                 subfolder="vae",
                 variant=variant,
-            ).to(self.pipe.device)
+            ).to(self.pipe.device, self.pipe.dtype)
             self.pipe.vae = torch.compile(
                 mode="reduce-overhead",
                 fullgraph=True,