adding enhanced garmet net

src/enhanced_garment_net.py

@@ -1,4 +1,3 @@
-"""
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -122,5 +121,3 @@ class EnhancedGarmentNetWithTimestep(nn.Module):
             combined_features.append(combined)
         
         return garment_out, combined_features
-
-"""        

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,7 +401,7 @@ class StableDiffusionXLInpaintPipeline(
         force_zeros_for_empty_prompt: bool = True,
     ):
         super().__init__()
-        # self.garment_net = EnhancedGarmentNetWithTimestep()
+        self.garment_net = EnhancedGarmentNetWithTimestep()
 
         self.register_modules(
             vae=vae,
@@ -1786,8 +1786,8 @@ class StableDiffusionXLInpaintPipeline(
                     added_cond_kwargs["image_embeds"] = image_embeds
                 print("Calling unet encoder for garment feature extraction")
                 # down,reference_features = self.UNet_Encoder(cloth,t, text_embeds_cloth,added_cond_kwargs= {"text_embeds": pooled_prompt_embeds_c, "time_ids": add_time_ids},return_dict=False)
-                down,reference_features = self.unet_encoder(cloth,t, text_embeds_cloth,return_dict=False)                
-                #garment_out, reference_features = self.garment_net(cloth, t, text_embeds_cloth)
+                #down,reference_features = self.unet_encoder(cloth,t, text_embeds_cloth,return_dict=False)                
+                garment_out, reference_features = self.garment_net(cloth, t, text_embeds_cloth)
                 print(type(reference_features))
                 print(reference_features)
                 reference_features = list(reference_features)