app.py

try:
    import spaces
    USING_SPACES = True
except ImportError:
    USING_SPACES = False

import sys
import os
import copy
import random
import numpy as np
import torchvision.transforms as T
import math
import peft
from peft import LoraConfig
from safetensors import safe_open
from omegaconf import OmegaConf
import gradio as gr
import torch
import diffusers
import transformers

os.environ["GRADIO_TEMP_DIR"] = ".gradio"

from omnitry.models.transformer_flux import FluxTransformer2DModel
from omnitry.pipelines.pipeline_flux_fill import FluxFillPipeline

# Setup device
if torch.cuda.is_available():
    device = torch.device('cuda:0')
else:
    device = torch.device('cpu')

HF_TOKEN = os.getenv("HF_TOKEN", "")
if HF_TOKEN:
    os.environ["HF_TOKEN"] = HF_TOKEN

weight_dtype = torch.bfloat16
args = OmegaConf.load('configs/omnitry_v1_unified.yaml')

# init model & pipeline
model_root = args.model_root
if not os.path.isdir(model_root):
    model_root = "black-forest-labs/FLUX.1-Fill-dev"

print(f"Loading base model from: {model_root} on device: {device}")

try:
    if os.path.isdir(model_root) and os.path.exists(os.path.join(model_root, 'transformer')):
        transformer = FluxTransformer2DModel.from_pretrained(
            f'{model_root}/transformer',
            torch_dtype=weight_dtype,
            low_cpu_mem_usage=True,
            token=HF_TOKEN or None,
        ).requires_grad_(False)
    else:
        try:
            transformer = FluxTransformer2DModel.from_pretrained(
                model_root,
                subfolder='transformer',
                torch_dtype=weight_dtype,
                low_cpu_mem_usage=True,
                token=HF_TOKEN or None,
            ).requires_grad_(False)
        except Exception as gated_err:
            print(f"Could not load gated repo '{model_root}', falling back to public mirror 'fuliucansheng/FLUX.1-Fill-dev-diffusers': {gated_err}")
            model_root = "fuliucansheng/FLUX.1-Fill-dev-diffusers"
            transformer = FluxTransformer2DModel.from_pretrained(
                model_root,
                subfolder='transformer',
                torch_dtype=weight_dtype,
                low_cpu_mem_usage=True,
                token=HF_TOKEN or None,
            ).requires_grad_(False)

    pipeline = FluxFillPipeline.from_pretrained(
        model_root,
        transformer=transformer.eval(),
        torch_dtype=weight_dtype,
        low_cpu_mem_usage=True,
        token=HF_TOKEN or None,
    )
except Exception as err:
    print("\n" + "="*80)
    print("ERROR LOADING MODEL CHECKPOINT:")
    print(err)
    print("="*80)
    print("\nNOTE: The base model 'black-forest-labs/FLUX.1-Fill-dev' is a gated Hugging Face repository.")
    print("Please ensure you have set HF_TOKEN in your environment or Hugging Face Space Secrets.")
    print("="*80 + "\n")
    sys.exit(1)

# VRAM saving
if torch.cuda.is_available():
    pipeline.enable_model_cpu_offload()

pipeline.vae.enable_tiling()

# 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'
    ]
)
print("[1/4] Adding vtryon_lora adapter...", flush=True)
transformer.add_adapter(lora_config, adapter_name='vtryon_lora')
print("[2/4] Adding garment_lora adapter...", flush=True)
transformer.add_adapter(lora_config, adapter_name='garment_lora')
print("[3/4] Loading LoRA weights from safetensors...", flush=True)

lora_path = args.lora_path
if not os.path.exists(lora_path):
    # Try downloading from HF hub if local safetensors missing
    from huggingface_hub import hf_hub_download
    print(f"Downloading omnitry lora weights from HF Hub...")
    lora_path = hf_hub_download(repo_id="Kunbyte/OmniTry", filename="omnitry_v1_unified.safetensors", token=HF_TOKEN or None)

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

print("[4/4] Patching LoRA forward functions...", flush=True)

# 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:
            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)

print("All setup complete. Initializing Gradio app...", flush=True)

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

def gpu_decorator(func):
    if USING_SPACES:
        return spaces.GPU(duration=120)(func)
    return func

@gpu_decorator
def generate(person_image, object_image, object_class, steps=20, guidance_scale=30, seed=-1, progress=gr.Progress(track_tqdm=True)):
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)
    seed_everything(seed)

    curr_device = torch.device('cuda:0') if torch.cuda.is_available() else device

    # 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=curr_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='cpu').manual_seed(seed),
        ).images[0]

    return img


with gr.Blocks() as demo:
    gr.Markdown('# 💎 OmniTry: Virtual Try-On for Jewelry, Hats & Accessories')
    with gr.Row():
        with gr.Column():
            person_image = gr.Image(type="pil", label="Person Image", height=600)
            run_button = gr.Button(value="Submit", variant='primary')

        with gr.Column():
            object_image = gr.Image(type="pil", label="Object Image (Jewelry / Hat / Watch / Accessory)", height=600)
            object_class = gr.Dropdown(label='Object Class', choices=list(args.object_map.keys()), value='earrings')

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

    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)

    ex_person = './demo_example/person_earrings.jpg' if os.path.exists('./demo_example/person_earrings.jpg') else None
    ex_object = './demo_example/object_earrings.jpg' if os.path.exists('./demo_example/object_earrings.jpg') else None
    if ex_person and ex_object:
        with gr.Row():
            gr.Examples(
                examples=[[ex_person, ex_object, 'earrings']],
                inputs=[person_image, object_image, object_class],
            )

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

if __name__ == '__main__':
    demo.queue().launch()