moving enhanced garment net initialization to app.py

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
 from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
 from src.unet_hacked_tryon import UNet2DConditionModel
+from src.enhanced_garment_net import EnhancedGarmentNetWithTimestep
 from transformers import (
     CLIPImageProcessor,
     CLIPVisionModelWithProjection,
@@ -51,6 +52,9 @@ unet = UNet2DConditionModel.from_pretrained(
     torch_dtype=torch.float16,
 )
 unet.requires_grad_(False)
+
+enhancedGarmentNet = EnhancedGarmentNetWithTimestep(dtype=torch.float16)
+
 tokenizer_one = AutoTokenizer.from_pretrained(
     base_path,
     subfolder="tokenizer",
@@ -92,6 +96,7 @@ UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
     torch_dtype=torch.float16,
 )
 
+
 parsing_model = Parsing(0)
 openpose_model = OpenPose(0)
 
@@ -122,6 +127,7 @@ pipe = TryonPipeline.from_pretrained(
         torch_dtype=torch.float16,
 )
 pipe.unet_encoder = UNet_Encoder
+pipe.garment_net = enhancedGarmentNet
 
 # Standard size of shein images
 #WIDTH = int(4160/5)
@@ -152,7 +158,8 @@ def start_tryon(human_img_dict,garm_img,garment_des, background_img, is_checked,
     openpose_model.preprocessor.body_estimation.model.to(device)
     pipe.to(device)
     pipe.unet_encoder.to(device)
-    
+    pipe.garment_net.to(device)
+
     human_img_orig = human_img_dict["background"].convert("RGB")   # ImageEditor
     #human_img_orig = human_img_dict.convert("RGB")     # Image
     

src/tryon_pipeline.py

@@ -57,7 +57,7 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 
 # Commenting out for now
-from src.enhanced_garment_net import EnhancedGarmentNetWithTimestep
+# from src.enhanced_garment_net import EnhancedGarmentNetWithTimestep
 
 
 
@@ -401,8 +401,8 @@ class StableDiffusionXLInpaintPipeline(
         force_zeros_for_empty_prompt: bool = True,
     ):
         super().__init__()
-        #self.garment_net = EnhancedGarmentNetWithTimestep()
-        self.garment_net = EnhancedGarmentNetWithTimestep().to(device=self._execution_device, dtype=self.unet.dtype)
+        # This is moved to app.py
+        #self.garment_net = EnhancedGarmentNetWithTimestep().to(device=self._execution_device, dtype=self.unet.dtype)