Spaces:
Runtime error
Runtime error
| # adapted from https://huggingface.co/spaces/HumanAIGC/OutfitAnyone/blob/main/app.py | |
| import os | |
| from os.path import join as opj | |
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| from torchvision import transforms | |
| from transformers import CLIPTextModel, CLIPTokenizer | |
| from transformers.models.clip.image_processing_clip import CLIPImageProcessor | |
| from diffusers import UniPCMultistepScheduler | |
| from diffusers import AutoencoderKL | |
| from diffusers import StableDiffusionPipeline | |
| from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker | |
| from stablegarment.models import GarmentEncoderModel,ControlNetModel | |
| from stablegarment.piplines import StableGarmentPipeline,StableGarmentControlNetPipeline | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch_dtype = torch.bfloat16 if device=="cpu" else torch.float16 | |
| height = 512 | |
| width = 384 | |
| base_model_path = "SG161222/Realistic_Vision_V4.0_noVAE" | |
| vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch_dtype,device=device) | |
| scheduler = UniPCMultistepScheduler.from_pretrained("runwayml/stable-diffusion-v1-5",subfolder="scheduler") | |
| pretrained_garment_encoder_path = "loooooong/StableGarment_text2img" | |
| garment_encoder = GarmentEncoderModel.from_pretrained(pretrained_garment_encoder_path,torch_dtype=torch_dtype,subfolder="garment_encoder") | |
| garment_encoder = garment_encoder.to(device=device,dtype=torch_dtype) | |
| pipeline_t2i = StableGarmentPipeline.from_pretrained(base_model_path, vae=vae, torch_dtype=torch_dtype, use_safetensors=True,).to(device=device) # variant="fp16" | |
| # pipeline = StableDiffusionPipeline.from_pretrained("SG161222/Realistic_Vision_V4.0_noVAE", vae=vae, torch_dtype=torch_dtype).to(device=device) | |
| pipeline_t2i.scheduler = scheduler | |
| pipeline_t2i.safety_checker = StableDiffusionSafetyChecker.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch_dtype, subfolder="safety_checker").to(device=device) | |
| pipeline_t2i.feature_extractor = CLIPImageProcessor.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch_dtype, subfolder="feature_extractor") | |
| pipeline_tryon = None | |
| ''' | |
| # not ready | |
| pretrained_model_path = "part_module_controlnet_imp2" | |
| controlnet = ControlNetModel.from_pretrained(pretrained_model_path,subfolder="controlnet") | |
| text_encoder = CLIPTextModel.from_pretrained(base_model_path, subfolder='text_encoder') | |
| tokenizer = CLIPTokenizer.from_pretrained(base_model_path, subfolder='tokenizer') | |
| pipeline_tryon = StableGarmentControlNetPipeline( | |
| vae, | |
| text_encoder, | |
| tokenizer, | |
| pipeline_t2i.unet, | |
| controlnet, | |
| scheduler, | |
| ).to(device=device,dtype=torch_dtype) | |
| ''' | |
| def prepare_controlnet_inputs(agn_mask_list,densepose_list): | |
| for i,agn_mask_img in enumerate(agn_mask_list): | |
| agn_mask_img = np.array(agn_mask_img.convert("L")) | |
| agn_mask_img = np.expand_dims(agn_mask_img, axis=-1) | |
| agn_mask_img = (agn_mask_img >= 128).astype(np.float32) # 0 or 1 | |
| agn_mask_list[i] = 1. - agn_mask_img | |
| densepose_list = [np.array(img)/255. for img in densepose_list] | |
| controlnet_inputs = [] | |
| for mask,pose in zip(agn_mask_list,densepose_list): | |
| controlnet_inputs.append(torch.tensor(np.concatenate([mask, pose], axis=-1)).permute(2,0,1)) | |
| controlnet_inputs = torch.stack(controlnet_inputs) | |
| return controlnet_inputs | |
| def tryon(prompt,init_image,garment_top,garment_down,): | |
| basename = os.path.splitext(os.path.basename(init_image))[0] | |
| image_agn = Image.open(opj(parse_dir,basename+"_agn.jpg")).resize((width,height)) | |
| image_agn_mask = Image.open(opj(parse_dir,basename+"_mask.png")).resize((width,height)) | |
| densepose_image = Image.open(opj(parse_dir,basename+"_densepose.png")).resize((width,height)) | |
| garment_top = Image.open(garment_top).resize((width,height)) | |
| garment_images = [garment_top,] | |
| prompt = [prompt,] | |
| cloth_prompt = ["",] | |
| controlnet_condition = prepare_controlnet_inputs([image_agn_mask],[densepose_image]).type(torch_dtype) | |
| images = pipeline_tryon(prompt, negative_prompt="",cloth_prompt=cloth_prompt, # negative_cloth_prompt = n_prompt, | |
| height=height,width=width,num_inference_steps=25,guidance_scale=1.5,eta=0.0, | |
| controlnet_condition=controlnet_condition,reference_image=garment_images, | |
| garment_encoder=garment_encoder,condition_extra=image_agn, | |
| generator=None,).images | |
| return images[0] | |
| def text2image(prompt,init_image,garment_top,garment_down,style_fidelity=1.): | |
| garment_top = Image.open(garment_top).resize((width,height)) | |
| garment_top = transforms.CenterCrop((height,width))(transforms.Resize(max(height, width))(garment_top)) | |
| # always enable classifier-free-guidance as it is related to garment | |
| cfg = 4 # if prompt else 0 | |
| garment_images = [garment_top,] | |
| prompt = [prompt,] | |
| cloth_prompt = ["",] | |
| n_prompt = "nsfw, unsaturated, abnormal, unnatural, artifact" | |
| negative_prompt = [n_prompt] | |
| images = pipeline_t2i(prompt,negative_prompt=negative_prompt,cloth_prompt=cloth_prompt,height=height,width=width, | |
| num_inference_steps=30,guidance_scale=cfg,num_images_per_prompt=1,style_fidelity=style_fidelity, | |
| garment_encoder=garment_encoder,garment_image=garment_images,).images | |
| return images[0] | |
| # def text2image(prompt,init_image,garment_top,garment_down,*args,**kwargs): | |
| # return pipeline(prompt).images[0] | |
| def infer(prompt,init_image,garment_top,garment_down,t2i_only,style_fidelity): | |
| if t2i_only: | |
| return text2image(prompt,init_image,garment_top,garment_down,style_fidelity) | |
| else: | |
| return tryon(prompt,init_image,garment_top,garment_down) | |
| init_state,prompt_state = None,"" | |
| t2i_only_state = True | |
| def set_mode(t2i_only,person_condition,prompt): | |
| global init_state, prompt_state, t2i_only_state | |
| t2i_only_state = not t2i_only_state | |
| init_state, prompt_state = person_condition or init_state, prompt_state or prompt | |
| if t2i_only: | |
| return [gr.Image(sources='clipboard', type="filepath", label="model",value=None, interactive=False), | |
| gr.Textbox(placeholder="", label="prompt(for t2i)", value=prompt_state, interactive=True), | |
| ] | |
| else: | |
| return [gr.Image(sources='clipboard', type="filepath", label="model",value=init_state, interactive=False), | |
| gr.Textbox(placeholder="", label="prompt(for t2i)", value="", interactive=False), | |
| ] | |
| def example_fn(inputs,): | |
| if t2i_only_state: | |
| return gr.Image(sources='clipboard', type="filepath", label="model", value=None, interactive=False) | |
| return gr.Image(sources='clipboard', type="filepath", label="model",value=inputs, interactive=False) | |
| gr.set_static_paths(paths=["assets/images/model"]) | |
| model_dir = opj(os.path.dirname(__file__), "assets/images/model") | |
| garment_dir = opj(os.path.dirname(__file__), "assets/images/garment") | |
| parse_dir = opj(os.path.dirname(__file__), "assets/images/image_parse") | |
| model = opj(model_dir, "13987_00.jpg") | |
| all_person = [opj(model_dir,fname) for fname in os.listdir(model_dir) if fname.endswith(".jpg")] | |
| with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px !important} ", ) as gradio_app: | |
| gr.Markdown("# StableGarment") | |
| gr.Markdown("Demo for [StableGarment: Garment-Centric Generation via Stable Diffusion](https://arxiv.org/abs/2403.10783).") | |
| gr.Markdown("*Running on cpu, so it is super slow. Feel free to duplicate the space or visit [StableGarment](https://github.com/logn-2024/StableGarment) for more info.*") | |
| with gr.Row(): | |
| with gr.Column(): | |
| init_image = gr.Image(sources='clipboard', type="filepath", label="model", value=None, interactive=False) | |
| example = gr.Examples(inputs=gr.Image(visible=False), #init_image, | |
| examples_per_page=4, | |
| examples=all_person, | |
| run_on_click=True, | |
| outputs=init_image, | |
| fn=example_fn,) | |
| with gr.Column(): | |
| with gr.Row(): | |
| images_top = [opj(garment_dir,fname) for fname in os.listdir(garment_dir) if fname.endswith(".jpg")] | |
| garment_top = gr.Image(sources='upload', type="filepath", label="top garment",value=images_top[0]) # ,interactive=False | |
| example_top = gr.Examples(inputs=garment_top, | |
| examples_per_page=4, | |
| examples=images_top) | |
| images_down = [] | |
| garment_down = gr.Image(sources='upload', type="filepath", label="lower garment",interactive=False, visible=False) | |
| example_down = gr.Examples(inputs=garment_down, | |
| examples_per_page=4, | |
| examples=images_down) | |
| prompt = gr.Textbox(placeholder="", label="prompt(for t2i)",) # interactive=False | |
| with gr.Row(): | |
| t2i_only = gr.Checkbox(label="t2i with garment", info="Only text and garment.", elem_id="t2i_switch", value=True, interactive=False,) | |
| run_button = gr.Button(value="Run") | |
| t2i_only.change(fn=set_mode,inputs=[t2i_only,init_image,prompt],outputs=[init_image,prompt,]) | |
| with gr.Accordion("advance options", open=False): | |
| gr.Markdown("Garment fidelity control(Tune down it to reduce white edge).") | |
| style_fidelity = gr.Slider(0, 1, value=1, label="fidelity(only for t2i)") # , info="" | |
| with gr.Column(): | |
| gallery = gr.Image() | |
| run_button.click(fn=infer, | |
| inputs=[ | |
| prompt, | |
| init_image, | |
| garment_top, | |
| garment_down, | |
| t2i_only, | |
| style_fidelity, | |
| ], | |
| outputs=[gallery],) | |
| gr.Markdown("We borrow some code from [OutfitAnyone](https://huggingface.co/spaces/HumanAIGC/OutfitAnyone), thanks. This demo is not safe for all audiences, which may reflect implicit bias and other defects of base model.") | |
| if __name__ == "__main__": | |
| gradio_app.launch() | |