Update app.py

app.py

@@ -2,75 +2,55 @@ import gradio as gr
 import torch
 from PIL import Image
 from diffusers import AutoPipelineForText2Image, DDIMScheduler
+from transformers import CLIPVisionModelWithProjection
 import numpy as np
-from torchvision import transforms
-import spaces
+import spaces  # Import ZeroGPU decorator
+
+# Load models and configure pipeline
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+    "h94/IP-Adapter",
+    subfolder="models/image_encoder",
+    torch_dtype=torch.float16,
+)
 
-# Initialize the pipeline
 pipeline = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
+    image_encoder=image_encoder,
 )
-
-# Configure the scheduler for the pipeline
 pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
 
-# Load IP adapter with specified weights and set the scale for each component
-pipeline.load_ip_adapter(
-    "h94/IP-Adapter",
-    subfolder="sdxl_models",
-    weight_name=[
-        "ip-adapter-plus_sdxl_vit-h.safetensors",
-        "ip-adapter-plus-face_sdxl_vit-h.safetensors"
-    ]
-)
+pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name=["ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus-face_sdxl_vit-h.safetensors"])
 pipeline.set_ip_adapter_scale([0.7, 0.5])
 
-# Define the desired size for the images
-desired_size = (1024, 1024)
+pipeline.enable_model_cpu_offload()
 
-@spaces.GPU
+@spaces.ZeroGPU  # Apply ZeroGPU decorator to the function
 def transform_image(face_image):
-    # Move the pipeline to the GPU inside the function
-    pipeline.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(0)
+    generator = torch.Generator(device="cpu").manual_seed(0)
 
-    # Process the input face image
+    # Check if the input is already a PIL Image
     if isinstance(face_image, Image.Image):
         processed_face_image = face_image
+    # If the input is a NumPy array, convert it to a PIL Image
     elif isinstance(face_image, np.ndarray):
         processed_face_image = Image.fromarray(face_image)
     else:
         raise ValueError("Unsupported image format")
-    
-    # Ensure the processed face image is in RGB format
-    processed_face_image = processed_face_image.convert('RGB')
-    
-    # Resize the face image to 1024x1024
-    processed_face_image = processed_face_image.resize(desired_size, Image.LANCZOS)
 
-    # Load the style image from the local path, resize it to 1024x1024, and convert to tensor
-    style_image_path = "examples/soyjak2.jpg"  # Ensure this path is correct
-    style_image = Image.open(style_image_path).resize(desired_size, Image.LANCZOS).convert('RGB')
-    style_image_tensor = transforms.ToTensor()(style_image).unsqueeze(0).to("cuda")
-    
-    # Convert the processed face image to tensor and move to GPU
-    processed_face_image_tensor = transforms.ToTensor()(processed_face_image).unsqueeze(0).to("cuda")
+    # Load the style image from the local path
+    style_image_path = "/content/soyjak2.jpeg"
+    style_image = Image.open(style_image_path)
 
-    # Perform the transformation using the configured pipeline
+    # Perform the transformation
     image = pipeline(
         prompt="soyjak",
-        ip_adapter_image=[style_image_tensor, processed_face_image_tensor],  # Ensure these are tensors
+        ip_adapter_image=[style_image, processed_face_image],
         negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
         num_inference_steps=30,
         generator=generator,
     ).images[0]
 
-    # Convert the tensor to a PIL Image to display it in Gradio
-    image = transforms.ToPILImage()(image.squeeze(0))
-    
-    # Move the pipeline back to CPU after processing to release GPU resources
-    pipeline.to("cpu")
     return image
 
 # Gradio interface setup
@@ -79,8 +59,8 @@ demo = gr.Interface(
     inputs=gr.Image(label="Upload your face image"),
     outputs=gr.Image(label="Your Soyjak"),
     title="InstaSoyjak - turn anyone into a Soyjak",
-    description="All you need to do is upload an image. Please use responsibly.",
+    description="All you need to do is upload an image. Please use responsibly. Please follow me on Twitter if you like this space: https://twitter.com/angrypenguinPNG. Idea from Yacine, please give him a follow: https://twitter.com/yacineMTB.",
 )
 
-demo.queue(max_size=20)
-demo.launch()
+demo.queue(max_size=20)  # Configures the queue with a maximum size of 20
+demo.launch()