From 4fb0ca5f750f5aae8068796b59705933100a4965 Mon Sep 17 00:00:00 2001 From: nftblackmagic Date: Tue, 26 Nov 2024 08:17:24 +0000 Subject: [PATCH] using lora version for spaces zeroGPU --- README.md | 10 ++- app.py | 24 ++++-- app_no_lora.py | 215 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 4 files changed, 241 insertions(+), 9 deletions(-) create mode 100644 app_no_lora.py diff --git a/README.md b/README.md index 9349a4f..32ec785 100644 --- a/README.md +++ b/README.md @@ -70,11 +70,17 @@ python tryon_inference.py \ --steps 30 ``` -Run the following command to start a gradio demo: +Run the following command to start a gradio demo with LoRA weights: ```bash python app.py ``` -Gradio demo: + +Run the following command to start a gradio demo without LoRA weights: +```bash +python app_no_lora.py +``` + + diff --git a/app.py b/app.py index 591da2f..2447929 100644 --- a/app.py +++ b/app.py @@ -37,16 +37,26 @@ def find_cuda(): device = torch.device('cuda') -print('Loading diffusion model ...') -transformer = FluxTransformer2DModel.from_pretrained( - "xiaozaa/catvton-flux-alpha", - torch_dtype=torch.bfloat16 +print("Start loading LoRA weights") +state_dict, network_alphas = FluxFillPipeline.lora_state_dict( + pretrained_model_name_or_path_or_dict="xiaozaa/catvton-flux-lora-alpha", ## The tryon Lora weights + weight_name="pytorch_lora_weights.safetensors", + return_alphas=True ) +is_correct_format = all("lora" in key or "dora_scale" in key for key in state_dict.keys()) +if not is_correct_format: + raise ValueError("Invalid LoRA checkpoint.") +print('Loading diffusion model ...') pipe = FluxFillPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - transformer=transformer, + "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16 ).to(device) +FluxFillPipeline.load_lora_into_transformer( + state_dict=state_dict, + network_alphas=network_alphas, + transformer=pipe.transformer, +) + print('Loading Finished!') @spaces.GPU @@ -99,7 +109,7 @@ def gradio_inference( with gr.Blocks() as demo: gr.Markdown(""" - # CATVTON FLUX Virtual Try-On Demo + # CATVTON FLUX Virtual Try-On Demo (by using LoRA weights) Upload a model image, draw a mask, and a garment image to generate virtual try-on results. [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/xiaozaa/catvton-flux-alpha) diff --git a/app_no_lora.py b/app_no_lora.py new file mode 100644 index 0000000..591da2f --- /dev/null +++ b/app_no_lora.py @@ -0,0 +1,215 @@ +import spaces + +import gradio as gr +from tryon_inference import run_inference +import os +import numpy as np +from PIL import Image +import tempfile +import torch +from diffusers import FluxTransformer2DModel, FluxFillPipeline + +import shutil + +def find_cuda(): + # Check if CUDA_HOME or CUDA_PATH environment variables are set + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + + if cuda_home and os.path.exists(cuda_home): + return cuda_home + + # Search for the nvcc executable in the system's PATH + nvcc_path = shutil.which('nvcc') + + if nvcc_path: + # Remove the 'bin/nvcc' part to get the CUDA installation path + cuda_path = os.path.dirname(os.path.dirname(nvcc_path)) + return cuda_path + + return None + +cuda_path = find_cuda() + +if cuda_path: + print(f"CUDA installation found at: {cuda_path}") +else: + print("CUDA installation not found") + +device = torch.device('cuda') + +print('Loading diffusion model ...') +transformer = FluxTransformer2DModel.from_pretrained( + "xiaozaa/catvton-flux-alpha", + torch_dtype=torch.bfloat16 +) +pipe = FluxFillPipeline.from_pretrained( + "black-forest-labs/FLUX.1-dev", + transformer=transformer, + torch_dtype=torch.bfloat16 +).to(device) +print('Loading Finished!') + +@spaces.GPU +def gradio_inference( + image_data, + garment, + num_steps=50, + guidance_scale=30.0, + seed=-1, + width=768, + height=1024 +): + """Wrapper function for Gradio interface""" + # 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") + temp_garment = os.path.join(tmp_dir, "garment.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) + garment.save(temp_garment) + + try: + # Run inference + _, tryon_result = run_inference( + pipe=pipe, + image_path=temp_image, + mask_path=temp_mask, + garment_path=temp_garment, + num_steps=num_steps, + guidance_scale=guidance_scale, + seed=seed, + size=(width, height) + ) + return tryon_result + except Exception as e: + raise gr.Error(f"Error during inference: {str(e)}") + +with gr.Blocks() as demo: + gr.Markdown(""" + # CATVTON FLUX Virtual Try-On Demo + Upload a model image, draw a mask, and a garment image to generate virtual try-on results. + + [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/xiaozaa/catvton-flux-alpha) + [![GitHub](https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white)](https://github.com/nftblackmagic/catvton-flux) + """) + + # gr.Video("example/github.mp4", label="Demo Video: How to use the tool") + + with gr.Column(): + 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/00057_00.jpg"], + ["./example/person/00067_00.jpg"], + ["./example/person/00069_00.jpg"], + ], + inputs=[image_input], + label="Person Images", + ) + with gr.Column(): + garment_input = gr.Image(label="Garment Image", type="pil", height=600, width=300) + gr.Examples( + examples=[ + ["./example/garment/04564_00.jpg"], + ["./example/garment/00055_00.jpg"], + ["./example/garment/00396_00.jpg"], + ["./example/garment/00067_00.jpg"], + ["./example/garment/00069_00.jpg"], + ], + inputs=[garment_input], + label="Garment Images", + ) + with gr.Column(): + tryon_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, + garment_input, + num_steps, + guidance_scale, + seed, + width, + height + ], + outputs=[tryon_output], + api_name="try-on" + ) + + +demo.launch() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index da050fd..398cd8b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,6 +8,7 @@ numpy==1.26.4 accelerate==1.1.1 sentencepiece==0.2.0 protobuf==5.27.3 +peft==0.13.2 huggingface-hub spaces git+https://github.com/huggingface/diffusers.git \ No newline at end of file