remove all optimization, add safety checker

app.py

@@ -12,12 +12,12 @@ 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.loaders import LoraLoaderMixin
-import intel_extension_for_pytorch as ipex
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 
 from stablegarment.models import GarmentEncoderModel,ControlNetModel
 from stablegarment.piplines import StableGarmentPipeline,StableGarmentControlNetPipeline
@@ -38,27 +38,8 @@ 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
-
-if device=="cpu":
-    # speed up for cpu
-    # to channels last
-    pipeline_t2i.unet = pipeline_t2i.unet.to(memory_format=torch.channels_last)
-    pipeline_t2i.vae = pipeline_t2i.vae.to(memory_format=torch.channels_last)
-    pipeline_t2i.text_encoder = pipeline_t2i.text_encoder.to(memory_format=torch.channels_last)
-    # pipeline_t2i.safety_checker = pipeline_t2i.safety_checker.to(memory_format=torch.channels_last)
-
-    # Create random input to enable JIT compilation
-    sample = torch.randn(2,4,64,48).type(torch_dtype)
-    timestep = torch.rand(1)*999
-    encoder_hidden_status = torch.randn(2,77,768).type(torch_dtype)
-    input_example = (sample, timestep, encoder_hidden_status)
-
-    # optimize with IPEX
-    pipeline_t2i.unet = ipex.optimize(pipeline_t2i.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
-    pipeline_t2i.vae = ipex.optimize(pipeline_t2i.vae.eval(), dtype=torch.bfloat16, inplace=True)
-    pipeline_t2i.text_encoder = ipex.optimize(pipeline_t2i.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
-    # pipeline_t2i.safety_checker = ipex.optimize(pipeline_t2i.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)
-
+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
 '''
@@ -77,7 +58,6 @@ pipeline_tryon = StableGarmentControlNetPipeline(
 ).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"))
@@ -101,7 +81,7 @@ def tryon(prompt,init_image,garment_top,garment_down,):
     garment_images = [garment_top,]
     prompt = [prompt,]
     cloth_prompt = ["",]
-    controlnet_condition = prepare_controlnet_inputs([image_agn_mask],[densepose_image])
+    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,
@@ -128,7 +108,7 @@ def text2image(prompt,init_image,garment_top,garment_down,style_fidelity=1.):
                     garment_encoder=garment_encoder,garment_image=garment_images,).images
     return images[0]
 
-# def text2image(prompt,init_image,garment_top,garment_down,):
+# 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):
@@ -166,6 +146,8 @@ 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)
@@ -207,6 +189,7 @@ with gr.Blocks(css = ".output-image, .input-image, .image-preview {height: 400px
                                     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()