Remove broken .so file from repo and LFS

34bd296 over 1 year ago

10.1 kB

import torch
from torchvision import transforms
from diffusers import DDPMScheduler, AutoencoderKL
from transformers import (
    CLIPImageProcessor,
    CLIPVisionModelWithProjection,
    CLIPTextModel,
    CLIPTextModelWithProjection,
    AutoTokenizer,
)
from src.tryon.src.tryon_pipeline import (
    StableDiffusionXLInpaintPipeline as TryonPipeline,
)
from src.tryon.src.unet_hacked_garmnet import (
    UNet2DConditionModel as UNet2DConditionModel_ref,
)
from src.tryon.src.unet_hacked_tryon import UNet2DConditionModel
from src.tryon.preprocess.humanparsing.run_parsing import Parsing
from src.tryon.preprocess.openpose.run_openpose import OpenPose
from src.tryon.utils_mask import get_mask_location
from torchvision.transforms.functional import to_pil_image
from src.tryon.detectron2.data.detection_utils import (
    convert_PIL_to_numpy,
    _apply_exif_orientation,
)
import torch
from PIL import Image
from src.tryon import apply_net


pipe = None
UNet_Encoder = None
parsing_model = None
openpose_model = None
tensor_transform = None


def load_model():
    global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform

    if pipe is not None:
        return {
            "pipe": pipe,
            "unet_encoder": UNet_Encoder,
            "parsing_model": parsing_model,
            "openpose_model": openpose_model,
            "tensor_transform": tensor_transform,
        }

    device = "cuda" if torch.cuda.is_available() else "cpu"
    base_path = "yisol/IDM-VTON"

    # Load sub-models
    unet = (
        UNet2DConditionModel.from_pretrained(
            base_path, subfolder="unet", torch_dtype=torch.float16
        )
        .eval()
        .requires_grad_(False)
    )

    tokenizer_one = AutoTokenizer.from_pretrained(
        base_path, subfolder="tokenizer", use_fast=False
    )
    tokenizer_two = AutoTokenizer.from_pretrained(
        base_path, subfolder="tokenizer_2", use_fast=False
    )

    noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")

    text_encoder_one = CLIPTextModel.from_pretrained(
        base_path, subfolder="text_encoder", torch_dtype=torch.float16
    )
    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
        base_path, subfolder="text_encoder_2", torch_dtype=torch.float16
    )
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        base_path, subfolder="image_encoder", torch_dtype=torch.float16
    )
    vae = AutoencoderKL.from_pretrained(
        base_path, subfolder="vae", torch_dtype=torch.float16
    )

    UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
        base_path, subfolder="unet_encoder", torch_dtype=torch.float16
    )
    UNet_Encoder.requires_grad_(False)

    parsing_model = Parsing(0)
    openpose_model = OpenPose(0)

    tensor_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )

    pipe = TryonPipeline.from_pretrained(
        base_path,
        unet=unet,
        vae=vae,
        feature_extractor=CLIPImageProcessor(),
        text_encoder=text_encoder_one,
        text_encoder_2=text_encoder_two,
        tokenizer=tokenizer_one,
        tokenizer_2=tokenizer_two,
        scheduler=noise_scheduler,
        image_encoder=image_encoder,
        torch_dtype=torch.float16,
    )
    pipe.unet_encoder = UNet_Encoder

    pipe.to(device)

    return {
        "pipe": pipe,
        "unet_encoder": UNet_Encoder,
        "parsing_model": parsing_model,
        "openpose_model": openpose_model,
        "tensor_transform": tensor_transform,
    }


def unload_model():
    global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform
    pipe = None
    UNet_Encoder = None
    parsing_model = None
    openpose_model = None
    tensor_transform = None
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def predict(
    # pipe,
    # UNet_Encoder,
    # parsing_model,
    # openpose_model,
    # tensor_transform,
    dict_img,
    garm_img,
    garment_des,
    is_checked,
    is_checked_crop,
    denoise_steps,
    seed,
    category,
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    openpose_model.preprocessor.body_estimation.model.to(device)
    pipe.unet_encoder.to(device)
    pipe.to(device)

    # Resize and convert
    garm_img = garm_img.convert("RGB").resize((768, 1024))
    human_img_orig = dict_img["background"].convert("RGB")
    orig_size = human_img_orig.size

    if is_checked_crop:
        width, height = human_img_orig.size
        target_width = int(min(width, height * (3 / 4)))
        target_height = int(min(height, width * (4 / 3)))
        left = (width - target_width) / 2
        top = (height - target_height) / 2
        right = (width + target_width) / 2
        bottom = (height + target_height) / 2
        cropped_img = human_img_orig.crop((left, top, right, bottom))
        crop_size = cropped_img.size
        human_img = cropped_img.resize((768, 1024))
    else:
        human_img = human_img_orig.resize((768, 1024))

    # Mask generation
    if is_checked:
        keypoints = openpose_model(human_img.resize((384, 512)))
        model_parse, _ = parsing_model(human_img.resize((384, 512)))
        mask, _ = get_mask_location("hd", category, model_parse, keypoints)
        mask = mask.resize((768, 1024))
    else:
        if dict_img["layers"][0] is not None:
            mask = dict_img["layers"][0].convert("RGB").resize((768, 1024))
            mask = Image.fromarray(
                (torch.tensor(mask) > 0).numpy().astype("uint8") * 255
            )
        else:
            mask = Image.new("RGB", (768, 1024), color="white")

    # Mask gray
    mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transform(human_img)
    mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)

    # Pose estimation
    human_np = convert_PIL_to_numpy(
        _apply_exif_orientation(human_img.resize((384, 512))), format="BGR"
    )
    args = apply_net.create_argument_parser().parse_args(
        (
            "show",
            "./tryon/configs/densepose_rcnn_R_50_FPN_s1x.yaml",
            "./tryon/ckpt/densepose/model_final_162be9.pkl",
            "dp_segm",
            "-v",
            "--opts",
            "MODEL.DEVICE",
            device,
        )
    )
    pose_img = args.func(args, human_np)
    pose_img = Image.fromarray(pose_img[:, :, ::-1]).resize((768, 1024))

    # Inference
    with torch.no_grad(), torch.cuda.amp.autocast():
        prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing {garment_des}"
        neg_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"

        prompt_embeds, neg_embeds, pooled_prompt_embeds, neg_pooled_prompt_embeds = (
            pipe.encode_prompt(
                prompt,
                num_images_per_prompt=1,
                do_classifier_free_guidance=True,
                negative_prompt=neg_prompt,
            )
        )

        garment_prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), a photo of {garment_des}"
        prompt_embeds_c, *_ = pipe.encode_prompt(
            [garment_prompt],
            num_images_per_prompt=1,
            do_classifier_free_guidance=False,
            negative_prompt=[neg_prompt],
        )

        pose_tensor = tensor_transform(pose_img).unsqueeze(0).to(device, torch.float16)
        garm_tensor = tensor_transform(garm_img).unsqueeze(0).to(device, torch.float16)
        generator = torch.Generator(device).manual_seed(seed) if seed != -1 else None

        images = pipe(
            prompt_embeds=prompt_embeds.to(device),
            negative_prompt_embeds=neg_embeds.to(device),
            pooled_prompt_embeds=pooled_prompt_embeds.to(device),
            negative_pooled_prompt_embeds=neg_pooled_prompt_embeds.to(device),
            num_inference_steps=denoise_steps,
            generator=generator,
            strength=1.0,
            pose_img=pose_tensor,
            text_embeds_cloth=prompt_embeds_c.to(device),
            cloth=garm_tensor,
            mask_image=mask,
            image=human_img,
            height=1024,
            width=768,
            ip_adapter_image=garm_img.resize((768, 1024)),
            guidance_scale=2.0,
        )[0]

    if is_checked_crop:
        return images[0].resize(crop_size), mask_gray.resize(crop_size)
    else:
        return images[0].resize(orig_size), mask_gray.resize(orig_size)


# from PIL import Image
# import torch
# from tryon_model import load_model, predict  # replace 'tryon_model' with your actual filename/module

# # Load the models
# pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform = load_model()

# # Load input images
# human_image = Image.open("path_to_human_image.jpg").convert("RGB")
# garment_image = Image.open("path_to_garment_image.jpg").convert("RGB")

# # Create mock dict_img
# dict_img = {
#     "background": human_image,
#     "layers": [None]  # or use a mask PIL.Image
# }

# # Set parameters
# garment_des = "a red floral dress with puff sleeves"
# is_checked = True             # Auto-generate mask
# is_checked_crop = True        # Crop input to 3:4 ratio
# denoise_steps = 30            # Diffusion steps
# seed = 42                     # Set -1 for random
# category = "upper"            # Use correct clothing type

# # Run prediction
# result_img, result_mask = predict(
#     pipe=pipe,
#     UNet_Encoder=UNet_Encoder,
#     parsing_model=parsing_model,
#     openpose_model=openpose_model,
#     tensor_transform=tensor_transform,
#     dict_img=dict_img,
#     garm_img=garment_image,
#     garment_des=garment_des,
#     is_checked=is_checked,
#     is_checked_crop=is_checked_crop,
#     denoise_steps=denoise_steps,
#     seed=seed,
#     category=category
# )

# # Save or show results
# result_img.save("tryon_output.png")
# result_mask.save("tryon_mask.png")
# print("Prediction complete. Saved output and mask.")

	import torch
	from torchvision import transforms
	from diffusers import DDPMScheduler, AutoencoderKL
	from transformers import (
	CLIPImageProcessor,
	CLIPVisionModelWithProjection,
	CLIPTextModel,
	CLIPTextModelWithProjection,
	AutoTokenizer,
	)
	from src.tryon.src.tryon_pipeline import (
	StableDiffusionXLInpaintPipeline as TryonPipeline,
	)
	from src.tryon.src.unet_hacked_garmnet import (
	UNet2DConditionModel as UNet2DConditionModel_ref,
	)
	from src.tryon.src.unet_hacked_tryon import UNet2DConditionModel
	from src.tryon.preprocess.humanparsing.run_parsing import Parsing
	from src.tryon.preprocess.openpose.run_openpose import OpenPose
	from src.tryon.utils_mask import get_mask_location
	from torchvision.transforms.functional import to_pil_image
	from src.tryon.detectron2.data.detection_utils import (
	convert_PIL_to_numpy,
	_apply_exif_orientation,
	)
	import torch
	from PIL import Image
	from src.tryon import apply_net


	pipe = None
	UNet_Encoder = None
	parsing_model = None
	openpose_model = None
	tensor_transform = None


	def load_model():
	global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform

	if pipe is not None:
	return {
	"pipe": pipe,
	"unet_encoder": UNet_Encoder,
	"parsing_model": parsing_model,
	"openpose_model": openpose_model,
	"tensor_transform": tensor_transform,
	}

	device = "cuda" if torch.cuda.is_available() else "cpu"
	base_path = "yisol/IDM-VTON"

	# Load sub-models
	unet = (
	UNet2DConditionModel.from_pretrained(
	base_path, subfolder="unet", torch_dtype=torch.float16
	)
	.eval()
	.requires_grad_(False)
	)

	tokenizer_one = AutoTokenizer.from_pretrained(
	base_path, subfolder="tokenizer", use_fast=False
	)
	tokenizer_two = AutoTokenizer.from_pretrained(
	base_path, subfolder="tokenizer_2", use_fast=False
	)

	noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")

	text_encoder_one = CLIPTextModel.from_pretrained(
	base_path, subfolder="text_encoder", torch_dtype=torch.float16
	)
	text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
	base_path, subfolder="text_encoder_2", torch_dtype=torch.float16
	)
	image_encoder = CLIPVisionModelWithProjection.from_pretrained(
	base_path, subfolder="image_encoder", torch_dtype=torch.float16
	)
	vae = AutoencoderKL.from_pretrained(
	base_path, subfolder="vae", torch_dtype=torch.float16
	)

	UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
	base_path, subfolder="unet_encoder", torch_dtype=torch.float16
	)
	UNet_Encoder.requires_grad_(False)

	parsing_model = Parsing(0)
	openpose_model = OpenPose(0)

	tensor_transform = transforms.Compose(
	[
	transforms.ToTensor(),
	transforms.Normalize([0.5], [0.5]),
	]
	)

	pipe = TryonPipeline.from_pretrained(
	base_path,
	unet=unet,
	vae=vae,
	feature_extractor=CLIPImageProcessor(),
	text_encoder=text_encoder_one,
	text_encoder_2=text_encoder_two,
	tokenizer=tokenizer_one,
	tokenizer_2=tokenizer_two,
	scheduler=noise_scheduler,
	image_encoder=image_encoder,
	torch_dtype=torch.float16,
	)
	pipe.unet_encoder = UNet_Encoder

	pipe.to(device)

	return {
	"pipe": pipe,
	"unet_encoder": UNet_Encoder,
	"parsing_model": parsing_model,
	"openpose_model": openpose_model,
	"tensor_transform": tensor_transform,
	}


	def unload_model():
	global pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform
	pipe = None
	UNet_Encoder = None
	parsing_model = None
	openpose_model = None
	tensor_transform = None
	if torch.cuda.is_available():
	torch.cuda.empty_cache()


	def predict(
	# pipe,
	# UNet_Encoder,
	# parsing_model,
	# openpose_model,
	# tensor_transform,
	dict_img,
	garm_img,
	garment_des,
	is_checked,
	is_checked_crop,
	denoise_steps,
	seed,
	category,
	):
	device = "cuda" if torch.cuda.is_available() else "cpu"
	openpose_model.preprocessor.body_estimation.model.to(device)
	pipe.unet_encoder.to(device)
	pipe.to(device)

	# Resize and convert
	garm_img = garm_img.convert("RGB").resize((768, 1024))
	human_img_orig = dict_img["background"].convert("RGB")
	orig_size = human_img_orig.size

	if is_checked_crop:
	width, height = human_img_orig.size
	target_width = int(min(width, height * (3 / 4)))
	target_height = int(min(height, width * (4 / 3)))
	left = (width - target_width) / 2
	top = (height - target_height) / 2
	right = (width + target_width) / 2
	bottom = (height + target_height) / 2
	cropped_img = human_img_orig.crop((left, top, right, bottom))
	crop_size = cropped_img.size
	human_img = cropped_img.resize((768, 1024))
	else:
	human_img = human_img_orig.resize((768, 1024))

	# Mask generation
	if is_checked:
	keypoints = openpose_model(human_img.resize((384, 512)))
	model_parse, _ = parsing_model(human_img.resize((384, 512)))
	mask, _ = get_mask_location("hd", category, model_parse, keypoints)
	mask = mask.resize((768, 1024))
	else:
	if dict_img["layers"][0] is not None:
	mask = dict_img["layers"][0].convert("RGB").resize((768, 1024))
	mask = Image.fromarray(
	(torch.tensor(mask) > 0).numpy().astype("uint8") * 255
	)
	else:
	mask = Image.new("RGB", (768, 1024), color="white")

	# Mask gray
	mask_gray = (1 - transforms.ToTensor()(mask)) * tensor_transform(human_img)
	mask_gray = to_pil_image((mask_gray + 1.0) / 2.0)

	# Pose estimation
	human_np = convert_PIL_to_numpy(
	_apply_exif_orientation(human_img.resize((384, 512))), format="BGR"
	)
	args = apply_net.create_argument_parser().parse_args(
	(
	"show",
	"./tryon/configs/densepose_rcnn_R_50_FPN_s1x.yaml",
	"./tryon/ckpt/densepose/model_final_162be9.pkl",
	"dp_segm",
	"-v",
	"--opts",
	"MODEL.DEVICE",
	device,
	)
	)
	pose_img = args.func(args, human_np)
	pose_img = Image.fromarray(pose_img[:, :, ::-1]).resize((768, 1024))

	# Inference
	with torch.no_grad(), torch.cuda.amp.autocast():
	prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), the model is wearing {garment_des}"
	neg_prompt = "monochrome, lowres, bad anatomy, worst quality, normal quality, low quality, blurry, jpeg artifacts, sketch"

	prompt_embeds, neg_embeds, pooled_prompt_embeds, neg_pooled_prompt_embeds = (
	pipe.encode_prompt(
	prompt,
	num_images_per_prompt=1,
	do_classifier_free_guidance=True,
	negative_prompt=neg_prompt,
	)
	)

	garment_prompt = f"((best quality, masterpiece, ultra-detailed, high quality photography, photo realistic)), a photo of {garment_des}"
	prompt_embeds_c, *_ = pipe.encode_prompt(
	[garment_prompt],
	num_images_per_prompt=1,
	do_classifier_free_guidance=False,
	negative_prompt=[neg_prompt],
	)

	pose_tensor = tensor_transform(pose_img).unsqueeze(0).to(device, torch.float16)
	garm_tensor = tensor_transform(garm_img).unsqueeze(0).to(device, torch.float16)
	generator = torch.Generator(device).manual_seed(seed) if seed != -1 else None

	images = pipe(
	prompt_embeds=prompt_embeds.to(device),
	negative_prompt_embeds=neg_embeds.to(device),
	pooled_prompt_embeds=pooled_prompt_embeds.to(device),
	negative_pooled_prompt_embeds=neg_pooled_prompt_embeds.to(device),
	num_inference_steps=denoise_steps,
	generator=generator,
	strength=1.0,
	pose_img=pose_tensor,
	text_embeds_cloth=prompt_embeds_c.to(device),
	cloth=garm_tensor,
	mask_image=mask,
	image=human_img,
	height=1024,
	width=768,
	ip_adapter_image=garm_img.resize((768, 1024)),
	guidance_scale=2.0,
	)[0]

	if is_checked_crop:
	return images[0].resize(crop_size), mask_gray.resize(crop_size)
	else:
	return images[0].resize(orig_size), mask_gray.resize(orig_size)


	# from PIL import Image
	# import torch
	# from tryon_model import load_model, predict # replace 'tryon_model' with your actual filename/module

	# # Load the models
	# pipe, UNet_Encoder, parsing_model, openpose_model, tensor_transform = load_model()

	# # Load input images
	# human_image = Image.open("path_to_human_image.jpg").convert("RGB")
	# garment_image = Image.open("path_to_garment_image.jpg").convert("RGB")

	# # Create mock dict_img
	# dict_img = {
	# "background": human_image,
	# "layers": [None] # or use a mask PIL.Image
	# }

	# # Set parameters
	# garment_des = "a red floral dress with puff sleeves"
	# is_checked = True # Auto-generate mask
	# is_checked_crop = True # Crop input to 3:4 ratio
	# denoise_steps = 30 # Diffusion steps
	# seed = 42 # Set -1 for random
	# category = "upper" # Use correct clothing type

	# # Run prediction
	# result_img, result_mask = predict(
	# pipe=pipe,
	# UNet_Encoder=UNet_Encoder,
	# parsing_model=parsing_model,
	# openpose_model=openpose_model,
	# tensor_transform=tensor_transform,
	# dict_img=dict_img,
	# garm_img=garment_image,
	# garment_des=garment_des,
	# is_checked=is_checked,
	# is_checked_crop=is_checked_crop,
	# denoise_steps=denoise_steps,
	# seed=seed,
	# category=category
	# )

	# # Save or show results
	# result_img.save("tryon_output.png")
	# result_mask.save("tryon_mask.png")
	# print("Prediction complete. Saved output and mask.")