Update app.py

app.py

@@ -2,11 +2,10 @@ 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 without specifying the device; this will be handled by the @spaces.GPU decorator
+# Initialize the pipeline
 pipeline = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch.float16
@@ -26,14 +25,13 @@ 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")
-
-# Define the desired size
+# Define the desired size for the images
 desired_size = (1024, 1024)
 
 @spaces.GPU
 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)
 
     # Process the input face image
@@ -60,7 +58,7 @@ def transform_image(face_image):
         generator=generator,
     ).images[0]
 
-    # Move the pipeline back to CPU after processing to release GPU resources
+    # Move the pipeline back to CPU after processing if necessary
     pipeline.to("cpu")
     return image
 
@@ -70,7 +68,7 @@ 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. 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.",
+    description="All you need to do is upload an image. Please use responsibly.",
 )
 
 demo.queue(max_size=20)