add pix2pix turbo

server/pipelines/pix2pix/model.py

@@ -0,0 +1,59 @@
+# https://github.com/GaParmar/img2img-turbo/blob/main/src/model.py
+from diffusers import DDPMScheduler
+
+
+def make_1step_sched():
+    noise_scheduler_1step = DDPMScheduler.from_pretrained(
+        "stabilityai/sd-turbo", subfolder="scheduler"
+    )
+    noise_scheduler_1step.set_timesteps(1, device="cuda")
+    noise_scheduler_1step.alphas_cumprod = noise_scheduler_1step.alphas_cumprod.cuda()
+    return noise_scheduler_1step
+
+
+def my_vae_encoder_fwd(self, sample):
+    sample = self.conv_in(sample)
+    l_blocks = []
+    # down
+    for down_block in self.down_blocks:
+        l_blocks.append(sample)
+        sample = down_block(sample)
+    # middle
+    sample = self.mid_block(sample)
+    sample = self.conv_norm_out(sample)
+    sample = self.conv_act(sample)
+    sample = self.conv_out(sample)
+    self.current_down_blocks = l_blocks
+    return sample
+
+
+def my_vae_decoder_fwd(self, sample, latent_embeds=None):
+    sample = self.conv_in(sample)
+    upscale_dtype = next(iter(self.up_blocks.parameters())).dtype
+    # middle
+    sample = self.mid_block(sample, latent_embeds)
+    sample = sample.to(upscale_dtype)
+    if not self.ignore_skip:
+        skip_convs = [
+            self.skip_conv_1,
+            self.skip_conv_2,
+            self.skip_conv_3,
+            self.skip_conv_4,
+        ]
+        # up
+        for idx, up_block in enumerate(self.up_blocks):
+            skip_in = skip_convs[idx](self.incoming_skip_acts[::-1][idx] * self.gamma)
+            # add skip
+            sample = sample + skip_in
+            sample = up_block(sample, latent_embeds)
+    else:
+        for idx, up_block in enumerate(self.up_blocks):
+            sample = up_block(sample, latent_embeds)
+    # post-process
+    if latent_embeds is None:
+        sample = self.conv_norm_out(sample)
+    else:
+        sample = self.conv_norm_out(sample, latent_embeds)
+    sample = self.conv_act(sample)
+    sample = self.conv_out(sample)
+    return sample

server/pipelines/pix2pix/pix2pix_turbo.py

@@ -0,0 +1,212 @@
+# https://github.com/GaParmar/img2img-turbo/blob/main/src/pix2pix_turbo.py
+import os
+import requests
+import sys
+import pdb
+import copy
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, PretrainedConfig, CLIPTextModel
+from diffusers import AutoencoderKL, UNet2DConditionModel, DDPMScheduler
+from diffusers.utils.peft_utils import set_weights_and_activate_adapters
+from peft import LoraConfig
+
+from pipelines.pix2pix.model import (
+    make_1step_sched,
+    my_vae_encoder_fwd,
+    my_vae_decoder_fwd,
+)
+
+
+class TwinConv(torch.nn.Module):
+    def __init__(self, convin_pretrained, convin_curr):
+        super(TwinConv, self).__init__()
+        self.conv_in_pretrained = copy.deepcopy(convin_pretrained)
+        self.conv_in_curr = copy.deepcopy(convin_curr)
+        self.r = None
+
+    def forward(self, x):
+        x1 = self.conv_in_pretrained(x).detach()
+        x2 = self.conv_in_curr(x)
+        return x1 * (1 - self.r) + x2 * (self.r)
+
+
+class Pix2Pix_Turbo(torch.nn.Module):
+    def __init__(self, name, ckpt_folder="checkpoints"):
+        super().__init__()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "stabilityai/sd-turbo", subfolder="tokenizer"
+        )
+        self.text_encoder = CLIPTextModel.from_pretrained(
+            "stabilityai/sd-turbo", subfolder="text_encoder"
+        ).cuda()
+        self.sched = make_1step_sched()
+
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-turbo", subfolder="vae")
+        unet = UNet2DConditionModel.from_pretrained(
+            "stabilityai/sd-turbo", subfolder="unet"
+        )
+
+        if name == "edge_to_image":
+            url = "https://www.cs.cmu.edu/~img2img-turbo/models/edge_to_image_loras.pkl"
+            os.makedirs(ckpt_folder, exist_ok=True)
+            outf = os.path.join(ckpt_folder, "edge_to_image_loras.pkl")
+            if not os.path.exists(outf):
+                print(f"Downloading checkpoint to {outf}")
+                response = requests.get(url, stream=True)
+                total_size_in_bytes = int(response.headers.get("content-length", 0))
+                block_size = 1024  # 1 Kibibyte
+                progress_bar = tqdm(
+                    total=total_size_in_bytes, unit="iB", unit_scale=True
+                )
+                with open(outf, "wb") as file:
+                    for data in response.iter_content(block_size):
+                        progress_bar.update(len(data))
+                        file.write(data)
+                progress_bar.close()
+                if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+                    print("ERROR, something went wrong")
+                print(f"Downloaded successfully to {outf}")
+            p_ckpt = outf
+            sd = torch.load(p_ckpt, map_location="cpu")
+            unet_lora_config = LoraConfig(
+                r=sd["rank_unet"],
+                init_lora_weights="gaussian",
+                target_modules=sd["unet_lora_target_modules"],
+            )
+
+        if name == "sketch_to_image_stochastic":
+            # download from url
+            url = "https://www.cs.cmu.edu/~img2img-turbo/models/sketch_to_image_stochastic_lora.pkl"
+            os.makedirs(ckpt_folder, exist_ok=True)
+            outf = os.path.join(ckpt_folder, "sketch_to_image_stochastic_lora.pkl")
+            if not os.path.exists(outf):
+                print(f"Downloading checkpoint to {outf}")
+                response = requests.get(url, stream=True)
+                total_size_in_bytes = int(response.headers.get("content-length", 0))
+                block_size = 1024  # 1 Kibibyte
+                progress_bar = tqdm(
+                    total=total_size_in_bytes, unit="iB", unit_scale=True
+                )
+                with open(outf, "wb") as file:
+                    for data in response.iter_content(block_size):
+                        progress_bar.update(len(data))
+                        file.write(data)
+                progress_bar.close()
+                if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+                    print("ERROR, something went wrong")
+                print(f"Downloaded successfully to {outf}")
+            p_ckpt = outf
+            sd = torch.load(p_ckpt, map_location="cpu")
+            unet_lora_config = LoraConfig(
+                r=sd["rank_unet"],
+                init_lora_weights="gaussian",
+                target_modules=sd["unet_lora_target_modules"],
+            )
+            convin_pretrained = copy.deepcopy(unet.conv_in)
+            unet.conv_in = TwinConv(convin_pretrained, unet.conv_in)
+
+        vae.encoder.forward = my_vae_encoder_fwd.__get__(
+            vae.encoder, vae.encoder.__class__
+        )
+        vae.decoder.forward = my_vae_decoder_fwd.__get__(
+            vae.decoder, vae.decoder.__class__
+        )
+        # add the skip connection convs
+        vae.decoder.skip_conv_1 = torch.nn.Conv2d(
+            512, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
+        ).cuda()
+        vae.decoder.skip_conv_2 = torch.nn.Conv2d(
+            256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
+        ).cuda()
+        vae.decoder.skip_conv_3 = torch.nn.Conv2d(
+            128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False
+        ).cuda()
+        vae.decoder.skip_conv_4 = torch.nn.Conv2d(
+            128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
+        ).cuda()
+        vae_lora_config = LoraConfig(
+            r=sd["rank_vae"],
+            init_lora_weights="gaussian",
+            target_modules=sd["vae_lora_target_modules"],
+        )
+        vae.decoder.ignore_skip = False
+        vae.add_adapter(vae_lora_config, adapter_name="vae_skip")
+        unet.add_adapter(unet_lora_config)
+        _sd_unet = unet.state_dict()
+        for k in sd["state_dict_unet"]:
+            _sd_unet[k] = sd["state_dict_unet"][k]
+        unet.load_state_dict(_sd_unet)
+        unet.enable_xformers_memory_efficient_attention()
+        _sd_vae = vae.state_dict()
+        for k in sd["state_dict_vae"]:
+            _sd_vae[k] = sd["state_dict_vae"][k]
+        vae.load_state_dict(_sd_vae)
+        unet.to("cuda")
+        vae.to("cuda")
+        unet.eval()
+        vae.eval()
+        self.unet, self.vae = unet, vae
+        self.vae.decoder.gamma = 1
+        self.timesteps = torch.tensor([999], device="cuda").long()
+        self.last_prompt = ""
+        self.caption_enc = None
+        self.device = "cuda"
+
+    def forward(self, c_t, prompt, deterministic=True, r=1.0, noise_map=1.0):
+        # encode the text prompt
+        if prompt != self.last_prompt:
+            caption_tokens = self.tokenizer(
+                prompt,
+                max_length=self.tokenizer.model_max_length,
+                padding="max_length",
+                truncation=True,
+                return_tensors="pt",
+            ).input_ids.cuda()
+            caption_enc = self.text_encoder(caption_tokens)[0]
+            self.caption_enc = caption_enc
+            self.last_prompt = prompt
+
+        if deterministic:
+            encoded_control = (
+                self.vae.encode(c_t).latent_dist.sample()
+                * self.vae.config.scaling_factor
+            )
+            model_pred = self.unet(
+                encoded_control,
+                self.timesteps,
+                encoder_hidden_states=self.caption_enc,
+            ).sample
+            x_denoised = self.sched.step(
+                model_pred, self.timesteps, encoded_control, return_dict=True
+            ).prev_sample
+            self.vae.decoder.incoming_skip_acts = self.vae.encoder.current_down_blocks
+            output_image = (
+                self.vae.decode(x_denoised / self.vae.config.scaling_factor).sample
+            ).clamp(-1, 1)
+        else:
+            # scale the lora weights based on the r value
+            self.unet.set_adapters(["default"], weights=[r])
+            set_weights_and_activate_adapters(self.vae, ["vae_skip"], [r])
+            encoded_control = (
+                self.vae.encode(c_t).latent_dist.sample()
+                * self.vae.config.scaling_factor
+            )
+            # combine the input and noise
+            unet_input = encoded_control * r + noise_map * (1 - r)
+            self.unet.conv_in.r = r
+            unet_output = self.unet(
+                unet_input,
+                self.timesteps,
+                encoder_hidden_states=self.caption_enc,
+            ).sample
+            self.unet.conv_in.r = None
+            x_denoised = self.sched.step(
+                unet_output, self.timesteps, unet_input, return_dict=True
+            ).prev_sample
+            self.vae.decoder.incoming_skip_acts = self.vae.encoder.current_down_blocks
+            self.vae.decoder.gamma = r
+            output_image = (
+                self.vae.decode(x_denoised / self.vae.config.scaling_factor).sample
+            ).clamp(-1, 1)
+        return output_image

server/pipelines/pix2pixTurbo.py

@@ -0,0 +1,127 @@
+import torch
+from torchvision import transforms
+
+from config import Args
+from pydantic import BaseModel, Field
+from PIL import Image
+from pipelines.pix2pix.pix2pix_turbo import Pix2Pix_Turbo
+from pipelines.utils.canny_gpu import SobelOperator
+
+
+default_prompt = "close-up photo of the joker"
+page_content = """
+<h1 class="text-3xl font-bold">Real-Time pix2pix_turbo</h1>
+<h3 class="text-xl font-bold">pix2pix turbo</h3>
+<p class="text-sm">
+    This demo showcases
+    <a
+    href="https://github.com/GaParmar/img2img-turbo"
+    target="_blank"
+    class="text-blue-500 underline hover:no-underline">One-Step Image Translation with Text-to-Image Models
+    </a>
+</p>
+"""
+
+
+class Pipeline:
+    class Info(BaseModel):
+        name: str = "img2img"
+        title: str = "Image-to-Image SDXL"
+        description: str = "Generates an image from a text prompt"
+        input_mode: str = "image"
+        page_content: str = page_content
+
+    class InputParams(BaseModel):
+        prompt: str = Field(
+            default_prompt,
+            title="Prompt",
+            field="textarea",
+            id="prompt",
+        )
+
+        width: int = Field(
+            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+        )
+        height: int = Field(
+            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+        )
+        strength: float = Field(
+            1.0,
+            min=0.01,
+            max=10.0,
+            step=0.001,
+            title="Strength",
+            field="range",
+            hide=True,
+            id="strength",
+        )
+        deterministic: bool = Field(
+            True,
+            hide=True,
+            title="Deterministic",
+            field="checkbox",
+            id="deterministic",
+        )
+        canny_low_threshold: float = Field(
+            0.31,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Canny Low Threshold",
+            field="range",
+            hide=True,
+            id="canny_low_threshold",
+        )
+        canny_high_threshold: float = Field(
+            0.125,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Canny High Threshold",
+            field="range",
+            hide=True,
+            id="canny_high_threshold",
+        )
+        debug_canny: bool = Field(
+            False,
+            title="Debug Canny",
+            field="checkbox",
+            hide=True,
+            id="debug_canny",
+        )
+
+    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
+        self.model = Pix2Pix_Turbo("edge_to_image")
+        self.canny_torch = SobelOperator(device=device)
+        self.device = device
+
+    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
+        # generator = torch.manual_seed(params.seed)
+        # pipe = self.pipes[params.base_model_id]
+
+        canny_pil, canny_tensor = self.canny_torch(
+            params.image,
+            params.canny_low_threshold,
+            params.canny_high_threshold,
+            output_type="pil,tensor",
+        )
+
+        with torch.no_grad():
+            canny_tensor = torch.cat((canny_tensor, canny_tensor, canny_tensor), dim=1)
+            output_image = self.model(
+                canny_tensor,
+                params.prompt,
+                params.deterministic,
+                params.strength,
+            )
+            output_pil = transforms.ToPILImage()(output_image[0].cpu() * 0.5 + 0.5)
+
+        result_image = output_pil
+        if params.debug_canny:
+            # paste control_image on top of result_image
+            w0, h0 = (200, 200)
+            control_image = canny_pil.resize((w0, h0))
+            w1, h1 = result_image.size
+            result_image.paste(control_image, (w1 - w0, h1 - h0))
+
+        return result_image

server/pipelines/utils/canny_gpu.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 from torchvision.transforms import ToTensor, ToPILImage
 from PIL import Image
 
+
 class SobelOperator(nn.Module):
     def __init__(self, device="cuda"):
         super(SobelOperator, self).__init__()
@@ -25,7 +26,13 @@ class SobelOperator(nn.Module):
         self.edge_conv_y.weight = nn.Parameter(sobel_kernel_y.view((1, 1, 3, 3)))
 
     @torch.no_grad()
-    def forward(self, image: Image.Image, low_threshold: float, high_threshold: float):
+    def forward(
+        self,
+        image: Image.Image,
+        low_threshold: float,
+        high_threshold: float,
+        output_type="pil",
+    ) -> Image.Image | torch.Tensor | tuple[Image.Image, torch.Tensor]:
         # Convert PIL image to PyTorch tensor
         image_gray = image.convert("L")
         image_tensor = ToTensor()(image_gray).unsqueeze(0).to(self.device)
@@ -41,4 +48,9 @@ class SobelOperator(nn.Module):
         edge[edge <= low_threshold] = 0.0
 
         # Convert the result back to a PIL image
-        return ToPILImage()(edge.squeeze(0).cpu())
+        if output_type == "pil":
+            return ToPILImage()(edge.squeeze(0).cpu())
+        elif output_type == "tensor":
+            return edge
+        elif output_type == "pil,tensor":
+            return ToPILImage()(edge.squeeze(0).cpu()), edge

server/requirements.txt

@@ -1,4 +1,4 @@
-diffusers==0.27.1
+diffusers==0.25.1
 transformers==4.36.2
 --extra-index-url https://download.pytorch.org/whl/cu121;
 torch==2.2.0
@@ -8,10 +8,11 @@ Pillow==10.2.0
 accelerate==0.25.0
 compel==2.0.2
 controlnet-aux==0.0.7
-peft==0.6.0
+peft==0.9.0
 xformers; sys_platform != 'darwin' or platform_machine != 'arm64'
 markdown2
 safetensors
 stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v1.0.4/stable_fast-1.0.4+torch220cu121-cp310-cp310-manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
 oneflow @ https://github.com/siliconflow/oneflow_releases/releases/download/community_cu121/oneflow-0.9.1.dev20240316+cu121-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform != 'darwin' or platform_machine != 'arm64'
-onediff @ git+https://github.com/siliconflow/onediff.git@main#egg=onediff ; sys_platform != 'darwin' or platform_machine != 'arm64'
+onediff @ git+https://github.com/siliconflow/onediff.git@main#egg=onediff ; sys_platform != 'darwin' or platform_machine != 'arm64'
+setuptools