Update app.py

app.py

@@ -40,11 +40,9 @@ class Args:
         self.seed = 42
         self.guidance_scale = 2.0
         self.mixed_precision = None
-        
-# Determine the device to be used for computations (CUDA if available)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-logger = get_logger(__name__, log_level="INFO")
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+
 
 def pil_to_tensor(images):
     images = np.array(images).astype(np.float32) / 255.0
@@ -56,44 +54,46 @@ def pil_to_tensor(images):
 args = Args()
 
 # Define the data type for model weights
-weight_dtype = torch.float32
+weight_dtype = torch.float16
 
 if args.seed is not None:
         set_seed(args.seed)
 
 
 # Load scheduler, tokenizer and models.
+
 noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
 vae = AutoencoderKL.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="vae",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
          )
 unet = UNet2DConditionModel.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="unet",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
         )
 image_encoder = CLIPVisionModelWithProjection.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="image_encoder",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
     )
 unet_encoder = UNet2DConditionModel_ref.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="unet_encoder",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
     )
 text_encoder_one = CLIPTextModel.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="text_encoder",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
         )
 text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="text_encoder_2",
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
     )
+
 tokenizer_one = AutoTokenizer.from_pretrained(
         args.pretrained_model_name_or_path,
         subfolder="tokenizer",
@@ -113,9 +113,8 @@ image_encoder.requires_grad_(False)
 unet_encoder.requires_grad_(False)
 text_encoder_one.requires_grad_(False)
 text_encoder_two.requires_grad_(False)
-unet_encoder.to(device, weight_dtype)
-unet.eval()
-unet_encoder.eval()
+unet_encoder.requires_grad_(False)
+
 
 pipe = TryonPipeline.from_pretrained(
             args.pretrained_model_name_or_path,
@@ -129,13 +128,11 @@ pipe = TryonPipeline.from_pretrained(
             scheduler = noise_scheduler,
             image_encoder=image_encoder,
             unet_encoder = unet_encoder,
-            torch_dtype=torch.float32,
-    ).to(device)
-# pipe.enable_sequential_cpu_offload()
-# pipe.enable_model_cpu_offload()
-# pipe.enable_vae_slicing()
-# Function to generate the image based on inputs
+            torch_dtype=torch.float16,
+    )
+
 def generate_virtual_try_on(person_image, cloth_image, mask_image, pose_image,cloth_des):
+    pipe.to(device)
     # Prepare the input images as tensors
     person_image = person_image.resize((args.width, args.height))
     cloth_image = cloth_image.resize((args.width, args.height))