Clean slate with jpg

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: IT Blender
+emoji: 💡
+colorFrom: gray
+colorTo: green
+sdk: gradio
+sdk_version: 5.33.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,177 @@
+import gradio as gr
+import torch
+import numpy as np
+import spaces
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+from diffusers import FluxPipeline
+from attention_processor import FluxBlendedAttnProcessor2_0
+
+from utils_sample import set_seed, resize_and_add_margin
+import torchvision.transforms.functional as F
+
+
+dtype = torch.bfloat16
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev", torch_dtype=dtype
+)
+pipe = pipe.to("cuda")
+
+
+@spaces.GPU
+def process_image_and_text(image, scale, seed, text):
+    set_seed(seed)
+    blended_attn_procs = {}
+    for name, _ in pipe.transformer.attn_processors.items():
+        if "single" in name:
+            blended_attn_procs[name] = FluxBlendedAttnProcessor2_0(3072, ba_scale=scale, num_ref=1)
+        else:
+            blended_attn_procs[name] = pipe.transformer.attn_processors[name]
+
+    pipe.transformer.set_attn_processor(blended_attn_procs)
+    pipe.to(dtype)
+
+    model_path = hf_hub_download(
+        repo_id="Wonwoong/IT-Blender",
+        filename="FLUX/it-blender.bin" # adjust the filename as needed
+    )
+    pretrained_blended_attn_weights = torch.load(model_path, map_location=pipe._execution_device)
+
+    key_changed_blended_attn_weights = {}
+    for key, value in pretrained_blended_attn_weights.items():
+        block_idx = int(key.split(".")[0]) - 21
+        k_or_v = key.split("_")[2]
+        changed_key = f'single_transformer_blocks.{block_idx}.attn.processor.blended_attention_{k_or_v}_proj.weight'
+        key_changed_blended_attn_weights[changed_key] = value.to(dtype)
+        
+    missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
+
+    image = Image.open(img_path).convert('RGB') 
+    image = resize_and_add_margin(image, target_size=512)
+    
+    image_list = [image]
+
+    out = pipe(
+        prompt=prompt,
+        height=512,
+        width=512,
+        max_sequence_length=256,
+        generator=torch.Generator().manual_seed(seed+10*j),
+        it_blender_image=image_list
+    ).images[0]
+
+    return out
+
+
+def get_samples():
+    sample_list = [
+        {
+            "image": "assets/0.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of a monster cartoon character, imaginative, creative, design",
+        },
+        {
+            "image": "assets/1.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of an owl cartoon character, imaginative, creative, design",
+        },
+        {
+            "image": "assets/2.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of a dragon, imaginative, creative, design",
+        },
+        {
+            "image": "assets/character1.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of a dragon, imaginative, creative, design",
+        },
+        {
+            "image": "assets/character2.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of a dragon, imaginative, creative, design",
+        },
+        {
+            "image": "assets/character3.jpg",
+            "scale": 0.6,
+            "seed": 42,
+            "text": "A photo of a dragon, imaginative, creative, design",
+        },
+        {
+            "image": "assets/graphic1.jpg",
+            "scale": 0.7,
+            "seed": 42,
+            "text": "A photo of a woman, imaginative, creative, design",
+        },
+        {
+            "image": "assets/product1.jpg",
+            "scale": 0.8,
+            "seed": 42,
+            "text": "A photo of a motorcycle, imaginative, creative, design",
+        }
+
+    ]
+    return [
+        [
+            Image.open(sample["image"]).resize((512, 512)),
+            sample["scale"],
+            sample["seed"],
+            sample["text"],
+        ]
+        for sample in sample_list
+    ]
+
+
+header = """
+# 💡 IT-Blender / FLUX
+<div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
+<a href="https://arxiv.org/abs/2411.15098"><img src="https://img.shields.io/badge/ArXiv-Paper-A42C25.svg" alt="arXiv"></a>
+<a href="https://imagineforme.github.io/"><img alt="Build" src="https://img.shields.io/badge/Project%20Page-ITBlender-yellow"></a> 
+<a href="https://github.com/WonwoongCho/IT-Blender"><img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub"></a>
+</div>
+"""
+
+
+def create_app():
+    
+    with gr.Blocks() as app:
+        gr.Markdown(header, elem_id="header")
+        with gr.Row(equal_height=False):
+            with gr.Column(variant="panel", elem_classes="inputPanel"):
+                original_image = gr.Image(
+                    type="pil", label="Condition Image", width=300, elem_id="input"
+                )
+
+                scale = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Guidance Scale")
+                seed = gr.Number(value=42, label="seed", precision=0)
+                text = gr.Textbox(lines=2, label="Text Prompt", elem_id="text")
+
+                submit_btn = gr.Button("Run", elem_id="submit_btn")
+
+            with gr.Column(variant="panel", elem_classes="outputPanel"):
+                output_image = gr.Image(type="pil", elem_id="output")
+
+        with gr.Row():
+            examples = gr.Examples(
+                examples=get_samples(),
+                inputs=[original_image, scale, seed, text],
+                label="Examples",
+            )
+        
+        submit_btn.click(
+            fn=process_image_and_text,
+            inputs=[original_image, scale, seed, text],
+            outputs=output_image,
+        )
+        
+    return app
+
+
+if __name__ == "__main__":
+    create_app().launch(debug=True, ssr_mode=False)

convert_png_to_jpg.py

@@ -0,0 +1,23 @@
+from PIL import Image
+import os
+
+def convert_png_to_jpg(input_folder, output_folder=None, quality=85):
+    if output_folder is None:
+        output_folder = input_folder
+
+    os.makedirs(output_folder, exist_ok=True)
+
+    for filename in os.listdir(input_folder):
+        if filename.lower().endswith(".png"):
+            png_path = os.path.join(input_folder, filename)
+            jpg_name = os.path.splitext(filename)[0] + ".jpg"
+            jpg_path = os.path.join(output_folder, jpg_name)
+
+            with Image.open(png_path) as img:
+                rgb_img = img.convert("RGB")  # remove alpha
+                rgb_img.save(jpg_path, "JPEG", quality=quality)
+
+            print(f"Converted: {filename} → {jpg_name}")
+
+# Example usage
+convert_png_to_jpg("assets/")

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+protobuf
+sentencepiece
+accelerate
+einops
+huggingface_hub
+git+https://github.com/WonwoongCho/diffusers@main#egg=diffusers

src/attention_processor.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+from typing import Callable, List, Optional, Tuple, Union
+
+class FluxBlendedAttnProcessor2_0(nn.Module):
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(self, hidden_dim, ba_scale=1.0, num_ref=1, temperature=1.2):
+        super().__init__()
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("FluxBlendedAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        
+        self.blended_attention_k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.blended_attention_v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
+        self.ba_scale = ba_scale
+        self.num_ref = num_ref
+        self.temperature = temperature # this is used only when num_ref > 1
+
+    def __call__(
+        self,
+        attn, #: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        is_negative_prompt: bool = False
+    ) -> torch.FloatTensor:
+        assert encoder_hidden_states is None, "It should be given as None because we are applying it-blender only to the single streams."
+        batch_size, _, _ = hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            normalized_query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        if image_rotary_emb is not None:
+            from diffusers.models.embeddings import apply_rotary_emb
+
+            query = apply_rotary_emb(normalized_query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+        
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+
+        # [noisy, clean]
+        chunk = batch_size//(1+self.num_ref)
+        ba_query = normalized_query[:chunk]  # noisy query
+
+        ba_key = self.blended_attention_k_proj(hidden_states[chunk:]) # clean key
+        ba_value = self.blended_attention_v_proj(hidden_states[chunk:]) # clean value
+
+        ba_key = ba_key.view(chunk, -1, attn.heads, head_dim).transpose(1, 2) # the -1 is gonna be multiplied by self.num_ref
+        ba_value = ba_value.view(chunk, -1, attn.heads, head_dim).transpose(1, 2)
+
+        ba_hidden_states = F.scaled_dot_product_attention(
+            ba_query, ba_key, ba_value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, scale=(1 / math.sqrt(ba_query.size(-1)))*self.temperature if self.num_ref > 1 else 1 / math.sqrt(ba_query.size(-1))
+        )
+
+        ba_hidden_states = ba_hidden_states.transpose(1, 2).reshape(chunk, -1, attn.heads * head_dim)
+        ba_hidden_states = ba_hidden_states.to(query.dtype)
+
+        zero_tensor_list = [torch.zeros_like(ba_hidden_states)]*self.num_ref
+        ba_hidden_states = torch.cat([ba_hidden_states]+zero_tensor_list, dim=0)    
+            
+        hidden_states = hidden_states + self.ba_scale * ba_hidden_states
+        
+        return hidden_states

src/utils_sample.py

@@ -0,0 +1,60 @@
+import random
+import numpy as np
+from PIL import Image
+import torch
+
+def set_seed(seed: int):
+    """
+    Set the seed for reproducibility across different libraries and devices.
+    
+    Args:
+        seed (int): The seed value to set.
+    """
+    # Set seed for Python's random module
+    random.seed(seed)
+    
+    # Set seed for NumPy
+    np.random.seed(seed)
+    
+    # Set seed for PyTorch CPU
+    torch.manual_seed(seed)
+    
+    # Set seed for PyTorch GPU (if using CUDA)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # For multi-GPU setups
+    
+    # Ensure deterministic results for CUDA operations (optional)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def resize_and_center_crop(image, target_size=512):
+    w, h = image.size
+    scale = target_size / min(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image_resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+
+    left = (new_w - target_size) // 2
+    top = (new_h - target_size) // 2
+    right = left + target_size
+    bottom = top + target_size
+    image_cropped = image_resized.crop((left, top, right, bottom))
+
+    return image_cropped
+
+
+def resize_and_add_margin(image, target_size=512, background_color=(255, 255, 255)):
+    w, h = image.size
+    scale = target_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image_resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+
+    new_image = Image.new("RGB", (target_size, target_size), background_color)
+
+    left = (target_size - new_w) // 2
+    top = (target_size - new_h) // 2
+    new_image.paste(image_resized, (left, top))
+
+    return new_image