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()