app.py

import gradio as gr
import spaces
import torch
import diffusers
import transformers
import copy
import random
import numpy as np
import torchvision.transforms as T
import math
import os
import peft
from peft import LoraConfig
from safetensors import safe_open
from omegaconf import OmegaConf
from omnitry.models.transformer_flux import FluxTransformer2DModel
from omnitry.pipelines.pipeline_flux_fill import FluxFillPipeline


from huggingface_hub import snapshot_download
snapshot_download(repo_id="Kunbyte/OmniTry", local_dir="./OmniTry")

device = torch.device('cuda:0')
weight_dtype = torch.bfloat16
args = OmegaConf.load('configs/omnitry_v1_unified.yaml')

# init model
transformer = FluxTransformer2DModel.from_pretrained('black-forest-labs/FLUX.1-Fill-dev', subfolder='transformer').requires_grad_(False).to(device, dtype=weight_dtype)
pipeline = FluxFillPipeline.from_pretrained(
    'black-forest-labs/FLUX.1-Fill-dev',
    transformer=transformer, 
    torch_dtype=weight_dtype
).to(device)


# insert LoRA
lora_config = LoraConfig(
    r=args.lora_rank,
    lora_alpha=args.lora_alpha,
    init_lora_weights="gaussian",
    target_modules=[
        'x_embedder',
        'attn.to_k', 'attn.to_q', 'attn.to_v', 'attn.to_out.0', 
        'attn.add_k_proj', 'attn.add_q_proj', 'attn.add_v_proj', 'attn.to_add_out', 
        'ff.net.0.proj', 'ff.net.2', 'ff_context.net.0.proj', 'ff_context.net.2', 
        'norm1_context.linear', 'norm1.linear', 'norm.linear', 'proj_mlp', 'proj_out'
    ]
)
transformer.add_adapter(lora_config, adapter_name='vtryon_lora')
transformer.add_adapter(lora_config, adapter_name='garment_lora')

with safe_open('OmniTry/omnitry_v1_unified.safetensors', framework="pt") as f:
    lora_weights = {k: f.get_tensor(k) for k in f.keys()}
    transformer.load_state_dict(lora_weights, strict=False)

# hack lora forward
def create_hacked_forward(module):

    def lora_forward(self, active_adapter, x, *args, **kwargs):
        result = self.base_layer(x, *args, **kwargs)
        if active_adapter is not None:
            torch_result_dtype = result.dtype
            lora_A = self.lora_A[active_adapter]
            lora_B = self.lora_B[active_adapter]
            dropout = self.lora_dropout[active_adapter]
            scaling = self.scaling[active_adapter]
            x = x.to(lora_A.weight.dtype)
            result = result + lora_B(lora_A(dropout(x))) * scaling
        return result
    
    def hacked_lora_forward(self, x, *args, **kwargs):
        return torch.cat((
            lora_forward(self, 'vtryon_lora', x[:1], *args, **kwargs),
            lora_forward(self, 'garment_lora', x[1:], *args, **kwargs),
        ), dim=0)
    
    return hacked_lora_forward.__get__(module, type(module))

for n, m in transformer.named_modules():
    if isinstance(m, peft.tuners.lora.layer.Linear):
        m.forward = create_hacked_forward(m)


def seed_everything(seed=0):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)



@spaces.GPU
def generate(person_image, object_image, object_class, steps, guidance_scale, seed):
    # set seed
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    seed_everything(seed)

    # resize model
    max_area = 1024 * 1024
    oW = person_image.width
    oH = person_image.height

    ratio = math.sqrt(max_area / (oW * oH))
    ratio = min(1, ratio)
    tW, tH = int(oW * ratio) // 16 * 16, int(oH * ratio) // 16 * 16
    transform = T.Compose([
        T.Resize((tH, tW)),
        T.ToTensor(),
    ])
    person_image = transform(person_image)

    # resize and padding garment
    ratio = min(tW / object_image.width, tH / object_image.height)
    transform = T.Compose([
        T.Resize((int(object_image.height * ratio), int(object_image.width * ratio))),
        T.ToTensor(),
    ])
    object_image_padded = torch.ones_like(person_image)
    object_image = transform(object_image)
    new_h, new_w = object_image.shape[1], object_image.shape[2]
    min_x = (tW - new_w) // 2
    min_y = (tH - new_h) // 2
    object_image_padded[:, min_y: min_y + new_h, min_x: min_x + new_w] = object_image

    # prepare prompts & conditions
    prompts = [args.object_map[object_class]] * 2
    img_cond = torch.stack([person_image, object_image_padded]).to(dtype=weight_dtype, device=device) 
    mask = torch.zeros_like(img_cond).to(img_cond)

    with torch.no_grad():
        img = pipeline(
            prompt=prompts,
            height=tH,
            width=tW,    
            img_cond=img_cond,
            mask=mask,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=torch.Generator(device).manual_seed(seed),
        ).images[0]

    return img


if __name__ == '__main__':

    with gr.Blocks() as demo:
        gr.Markdown('# Demo of OmniTry')
        with gr.Row():
            with gr.Column():
                person_image = gr.Image(type="pil", label="Person Image", height=800)
                run_button = gr.Button(value="Submit", variant='primary')

            with gr.Column():
                object_image = gr.Image(type="pil", label="Object Image", height=800)
                object_class = gr.Dropdown(label='Object Class', choices=args.object_map.keys())

            with gr.Column():
                image_out = gr.Image(type="pil", label="Output", height=800)

        with gr.Accordion("Advanced ⚙️", open=False):
            guidance_scale = gr.Slider(label="Guidance scale", minimum=1, maximum=50, value=30, step=0.1)
            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
            seed = gr.Number(label="Seed", value=-1, precision=0)

        with gr.Row():
            gr.Examples(
                examples=[
                    [
                        './demo_example/person_top_cloth.jpg',
                        './demo_example/object_top_cloth.jpg', 
                        'top clothes',
                    ],
                    [
                        './demo_example/person_bottom_cloth.jpg',
                        './demo_example/object_bottom_cloth.jpg', 
                        'bottom clothes',
                    ],
                    [
                        './demo_example/person_dress.jpg',
                        './demo_example/object_dress.jpg', 
                        'dress',
                    ],
                    [
                        './demo_example/person_shoes.jpg',
                        './demo_example/object_shoes.jpg', 
                        'shoe',
                    ],
                    [
                        './demo_example/person_earrings.jpg',
                        './demo_example/object_earrings.jpg', 
                        'earrings',
                    ],
                    [
                        './demo_example/person_bracelet.jpg',
                        './demo_example/object_bracelet.jpg', 
                        'bracelet',
                    ],
                    [
                        './demo_example/person_necklace.jpg',
                        './demo_example/object_necklace.jpg', 
                        'necklace',
                    ],
                    [
                        './demo_example/person_ring.jpg',
                        './demo_example/object_ring.jpg', 
                        'ring',
                    ],
                    [
                        './demo_example/person_sunglasses.jpg',
                        './demo_example/object_sunglasses.jpg', 
                        'sunglasses',
                    ],
                    [
                        './demo_example/person_glasses.jpg',
                        './demo_example/object_glasses.jpg', 
                        'glasses',
                    ],
                    [
                        './demo_example/person_belt.jpg',
                        './demo_example/object_belt.jpg', 
                        'belt',
                    ],
                    [
                        './demo_example/person_bag.jpg',
                        './demo_example/object_bag.jpg', 
                        'bag',
                    ],
                    [
                        './demo_example/person_hat.jpg',
                        './demo_example/object_hat.jpg', 
                        'hat',
                    ],
                    [
                        './demo_example/person_tie.jpg',
                        './demo_example/object_tie.jpg', 
                        'tie',
                    ],
                    [
                        './demo_example/person_bowtie.jpg',
                        './demo_example/object_bowtie.jpg', 
                        'bow tie',
                    ],
                ],

                inputs=[person_image, object_image, object_class],
                examples_per_page=100
            )
            
        run_button.click(generate, inputs=[person_image, object_image, object_class, steps, guidance_scale, seed], outputs=[image_out])

    demo.launch(server_name="0.0.0.0")