try off version

.gitignore

@@ -0,0 +1,58 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+dist/
+build/
+*.egg-info/
+
+# Virtual environments
+venv/
+env/
+.env/
+.venv/
+
+# IDE specific files
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+coverage.xml
+*.cover
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# Local development settings
+.env
+.env.local
+
+# Logs
+*.log
+
+# Database files
+*.db
+*.sqlite3
+
+# OS generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db 
+
+# Gradio cache
+.gradio/example/github.mp4
+
+aws/
+checkpoints/

README.md

@@ -1,14 +1,17 @@
 ---
-title: Cat Try Off Flux
-emoji: 👀
+title: cat-tryoff-flux
+emoji: 🖥️
 colorFrom: yellow
-colorTo: purple
+colorTo: pink
 sdk: gradio
-sdk_version: 5.8.0
+sdk_version: 5.0.1
 app_file: app.py
 pinned: false
-license: cc-by-nc-4.0
-short_description: Extract and reconstruct the front view of clothing
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# cat-tryoff-flux
+
+CAT-Tryoff-Flux is an advanced tryoff model. This model can extract and reconstruct the front view of clothing items from images of people wearing them. It used the same method of (CATVTON-FLUX)[https://huggingface.co/xiaozaa/catvton-flux-alpha].
+
+The github repo is [here](https://github.com/nftblackmagic/catvton-flux).

app.py

@@ -1,7 +1,189 @@
+import spaces
+
 import gradio as gr
+from tryoff_inference import run_inference
+import os
+import numpy as np
+from PIL import Image
+import tempfile
+import torch
+from diffusers import FluxTransformer2DModel, FluxFillPipeline
+import subprocess
+
+subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print('Loading diffusion model ...')
+transformer = FluxTransformer2DModel.from_pretrained(
+    "xiaozaa/cat-tryoff-flux", 
+    torch_dtype=dtype
+)
+pipe = FluxFillPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    transformer=transformer,
+    torch_dtype=dtype
+).to(device)
+print('Loading Finished!')
+
+@spaces.GPU(duration=120)
+def gradio_inference(
+    image_data, 
+    garment, 
+    num_steps=50, 
+    guidance_scale=30.0, 
+    seed=-1, 
+    width=768,
+    height=1024
+):
+    """Wrapper function for Gradio interface"""
+    # Check if mask has been drawn
+    if image_data is None or "layers" not in image_data or not image_data["layers"]:
+        raise gr.Error("Please draw a mask over the clothing area before generating!")
+    
+    # Check if mask is empty (all black)
+    mask = image_data["layers"][0]
+    mask_array = np.array(mask)
+    if np.all(mask_array < 10):
+        raise gr.Error("The mask is empty! Please draw over the clothing area you want to replace.")
+
+    # Use temporary directory
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        # Save inputs to temp directory
+        temp_image = os.path.join(tmp_dir, "image.png")
+        temp_mask = os.path.join(tmp_dir, "mask.png")
+        
+        # Extract image and mask from ImageEditor data
+        image = image_data["background"]
+        mask = image_data["layers"][0]  # First layer contains the mask
+        
+        # Convert to numpy array and process mask
+        mask_array = np.array(mask)
+        is_black = np.all(mask_array < 10, axis=2)
+        mask = Image.fromarray(((~is_black) * 255).astype(np.uint8))
+        
+        # Save files to temp directory
+        image.save(temp_image)
+        mask.save(temp_mask)
+        
+        try:
+            # Run inference
+            garment_result, _ = run_inference(
+                pipe=pipe,
+                image_path=temp_image,
+                mask_path=temp_mask,
+                num_steps=num_steps,
+                guidance_scale=guidance_scale,
+                seed=seed,
+                size=(width, height)
+            )
+            return garment_result
+        except Exception as e:
+            raise gr.Error(f"Error during inference: {str(e)}")
+
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    # CAT-TRYOFF-FLUX Virtual Try-Off Demo
+    Upload a model image, draw a mask, and a garment image to generate virtual try-off results.
+    
+    """)
+    
+    # gr.Video("example/github.mp4", label="Demo Video: How to use the tool")
+    
+    with gr.Column():
+        gr.Markdown("""
+        ### ⚠️ Important:
+        1. Choose a model image or upload your own
+        2. Use the Pen tool to draw a mask over the clothing area you want to restore
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.ImageMask(
+                    label="Model Image (Click 'Edit' and draw mask over the clothing area)",
+                    type="pil",
+                    height=600,
+                    width=300
+                )
+                gr.Examples(
+                    examples=[
+                        ["./example/person/00008_00.jpg"],
+                        ["./example/person/00055_00.jpg"],
+                        ["./example/person/00064_00.jpg"],
+                        ["./example/person/00067_00.jpg"],
+                        ["./example/person/00069_00.jpg"],
+                    ],
+                    inputs=[image_input],
+                    label="Person Images",
+                )
+            with gr.Column():
+                garment_output = gr.Image(label="Try-On Result", height=600, width=300)
+            
+        with gr.Row():
+            num_steps = gr.Slider(
+                minimum=1, 
+                maximum=100, 
+                value=30, 
+                step=1, 
+                label="Number of Steps"
+            )
+            guidance_scale = gr.Slider(
+                minimum=1.0, 
+                maximum=50.0, 
+                value=30.0, 
+                step=0.5, 
+                label="Guidance Scale"
+            )
+            seed = gr.Slider(
+                minimum=-1,
+                maximum=2147483647,
+                step=1,
+                value=-1,
+                label="Seed (-1 for random)"
+            )
+            width = gr.Slider(
+                minimum=256,
+                maximum=1024,
+                step=64,
+                value=768,
+                label="Width"
+            )
+            height = gr.Slider(
+                minimum=256,
+                maximum=1024,
+                step=64,
+                value=1024,
+                label="Height"
+            )
+            
+        
+        submit_btn = gr.Button("Generate Try-On", variant="primary")
+        
+            
+    with gr.Row():
+        gr.Markdown("""
+        ### Notes:
+        - The model is trained on VITON-HD dataset. It focuses on the woman upper body try-on generation.
+        - The mask should indicate the region where the garment will be placed.
+        - The garment image should be on a clean background.
+        - The model is not perfect. It may generate some artifacts.
+        - The model is slow. Please be patient.
+        - The model is just for research purpose.
+        """)
+    
+    submit_btn.click(
+        fn=gradio_inference,
+        inputs=[
+            image_input,
+            num_steps,
+            guidance_scale,
+            seed,
+            width,
+            height
+        ],
+        outputs=[garment_output],
+        api_name="try-off"
+    )
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

requirements.txt

@@ -0,0 +1,14 @@
+accelerate
+git+https://github.com/huggingface/diffusers.git
+gradio==5.6.0
+gradio_client==1.4.3
+torch==2.4.0
+torchvision==0.19.0
+tqdm==4.66.5
+transformers==4.43.3
+numpy==1.26.4
+sentencepiece
+peft==0.13.2
+huggingface-hub
+spaces
+protobuf

tryoff.sh

@@ -0,0 +1,7 @@
+python tryoff_inference.py \
+--image ./example/person/00069_00.jpg \
+--mask ./example/person/00069_00_mask.png \
+--seed 41 \
+--output_tryon test_original.png \
+--output_garment restored_garment6.png \
+--steps 30

tryoff_inference.py

@@ -0,0 +1,117 @@
+import argparse
+import torch
+from diffusers.utils import load_image, check_min_version
+from diffusers import FluxPriorReduxPipeline, FluxFillPipeline
+from diffusers import FluxTransformer2DModel
+import numpy as np
+from torchvision import transforms
+
+def run_inference(
+    image_path,
+    mask_path,
+    size=(576, 768),
+    num_steps=50,
+    guidance_scale=30,
+    seed=42,
+    pipe=None
+):
+    # Build pipeline
+    if pipe is None:
+        transformer = FluxTransformer2DModel.from_pretrained(
+            "xiaozaa/cat-tryoff-flux", 
+            torch_dtype=torch.bfloat16
+        )
+        pipe = FluxFillPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev",
+            transformer=transformer,
+            torch_dtype=torch.bfloat16
+        ).to("cuda")
+    else:
+        pipe.to("cuda")
+
+    pipe.transformer.to(torch.bfloat16)
+
+    # Add transform
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])  # For RGB images
+    ])
+    mask_transform = transforms.Compose([
+        transforms.ToTensor()
+    ])
+
+    # Load and process images
+    # print("image_path", image_path)
+    image = load_image(image_path).convert("RGB").resize(size)
+    mask = load_image(mask_path).convert("RGB").resize(size)
+
+    # Transform images using the new preprocessing
+    image_tensor = transform(image)
+    mask_tensor = mask_transform(mask)[:1]  # Take only first channel
+    garment_tensor = torch.zeros_like(image_tensor)
+    image_tensor = image_tensor * mask_tensor
+
+    # Create concatenated images
+    inpaint_image = torch.cat([garment_tensor, image_tensor], dim=2)  # Concatenate along width
+    garment_mask = torch.zeros_like(mask_tensor)
+    extended_mask = torch.cat([1 - garment_mask, garment_mask], dim=2)
+
+    prompt = f"The pair of images highlights a clothing and its styling on a model, high resolution, 4K, 8K; " \
+            f"[IMAGE1] Detailed product shot of a clothing" \
+            f"[IMAGE2] The same cloth is worn by a model in a lifestyle setting."
+
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+
+    result = pipe(
+        height=size[1],
+        width=size[0] * 2,
+        image=inpaint_image,
+        mask_image=extended_mask,
+        num_inference_steps=num_steps,
+        generator=generator,
+        max_sequence_length=512,
+        guidance_scale=guidance_scale,
+        prompt=prompt,
+    ).images[0]
+
+    # Split and save results
+    width = size[0]
+    garment_result = result.crop((0, 0, width, size[1]))
+    tryon_result = result.crop((width, 0, width * 2, size[1]))
+    
+    return garment_result, tryon_result
+
+def main():
+    parser = argparse.ArgumentParser(description='Run FLUX virtual try-on inference')
+    parser.add_argument('--image', required=True, help='Path to the model image')
+    parser.add_argument('--mask', required=True, help='Path to the agnostic mask')
+    parser.add_argument('--output_garment', default='flux_inpaint_garment.png', help='Output path for garment result')
+    parser.add_argument('--output_tryon', default='flux_inpaint_tryon.png', help='Output path for try-on result')
+    parser.add_argument('--steps', type=int, default=50, help='Number of inference steps')
+    parser.add_argument('--guidance_scale', type=float, default=30, help='Guidance scale')
+    parser.add_argument('--seed', type=int, default=0, help='Random seed')
+    parser.add_argument('--width', type=int, default=576, help='Width')
+    parser.add_argument('--height', type=int, default=768, help='Height')
+    
+    args = parser.parse_args()
+    
+    check_min_version("0.30.2")
+    
+    garment_result, tryon_result = run_inference(
+        image_path=args.image,
+        mask_path=args.mask,
+        num_steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        seed=args.seed,
+        size=(args.width, args.height)
+    )
+    output_tryon_path=args.output_tryon
+    output_garment_path=args.output_garment
+    
+    tryon_result.save(output_tryon_path)
+    garment_result.save(output_garment_path)
+    
+    print("Successfully saved garment and try-on images")
+
+if __name__ == "__main__":
+    main()