Update app.py

app.py

@@ -45,22 +45,26 @@ def transform_image(face_image):
     # Resize the face image
     processed_face_image = processed_face_image.resize(desired_size, Image.LANCZOS)
 
-    # Load and resize the style image from the local path
-    style_image_path = "examples/soyjak2.jpeg"
-    style_image = Image.open(style_image_path).resize(desired_size, Image.LANCZOS)
+    # Convert PIL images to PyTorch tensors
+    processed_face_tensor = transforms.ToTensor()(processed_face_image).unsqueeze(0).to("cuda")
+    style_image_tensor = transforms.ToTensor()(style_image).unsqueeze(0).to("cuda")
+
+    # Ensure tensors are the correct shape (C, H, W)
+    if processed_face_tensor.shape[1:] != (3, 1280, 1280):
+        raise ValueError(f"Face image tensor shape is {processed_face_tensor.shape}, but expected shape is (3, 1280, 1280)")
 
     # Perform the transformation using the configured pipeline
     image = pipeline(
         prompt="soyjak",
-        ip_adapter_image=[style_image, processed_face_image],
+        ip_adapter_image=[style_image_tensor, processed_face_tensor],
         negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
         num_inference_steps=30,
         generator=generator,
     ).images[0]
 
-    # Move the pipeline back to CPU after processing if necessary
+    # Move the pipeline back to CPU after processing to release GPU resources
     pipeline.to("cpu")
-    return image
+    return transforms.ToPILImage()(image.squeeze(0))
 
 # Gradio interface setup
 demo = gr.Interface(