Update app.py

app.py

@@ -2,14 +2,13 @@ import gradio as gr
 import torch
 from PIL import Image
 from diffusers import AutoPipelineForText2Image, DDIMScheduler
-from transformers import CLIPVisionModelWithProjection
 import numpy as np
+import spaces  # Make sure to import spaces
 
-# Initialize the pipeline with GPU support
+# Initialize the pipeline without specifying the device; this will be handled by the @spaces.GPU decorator
 pipeline = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    device="cuda",  # Use GPU device if available
+    torch_dtype=torch.float16
 )
 
 # Configure the scheduler for the pipeline
@@ -26,10 +25,11 @@ pipeline.load_ip_adapter(
 )
 pipeline.set_ip_adapter_scale([0.7, 0.5])
 
-# Ensure the model and its components are moved to GPU
-pipeline.to("cuda")
-
+# Decorate the transform_image function to run on GPU
+@spaces.GPU
 def transform_image(face_image):
+    # Move the pipeline to GPU inside the function
+    pipeline.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(0)
 
     # Process the input face image
@@ -49,10 +49,12 @@ def transform_image(face_image):
         prompt="soyjak",
         ip_adapter_image=[style_image, processed_face_image],
         negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
-        num_inference_steps=30,
+        num_inference_steps=50,
         generator=generator,
     ).images[0]
 
+    # Move the pipeline back to CPU after processing to release GPU resources
+    pipeline.to("cpu")
     return image
 
 # Gradio interface setup