app.py

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  # Ensure this is available in your environment

# Initialize a zero tensor for demonstration purposes
zero = torch.Tensor([0]).cuda()
print(zero.device)  # Should output 'cuda:0' if a GPU is available

@spaces.GPU  # Decorate the function to run on GPU
def transform_image(face_image):
    print(zero.device)  # Check the device inside the function, should be 'cuda:0'
    
    generator = torch.Generator(device="cuda").manual_seed(0)  # Use GPU device if available

    # Process the input face image
    if isinstance(face_image, Image.Image):
        processed_face_image = face_image
    elif isinstance(face_image, np.ndarray):
        processed_face_image = Image.fromarray(face_image)
    else:
        raise ValueError("Unsupported image format")

    # 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 GPU
    image = pipeline(
        prompt="soyjak",
        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]

    return image

# Load models and configure pipeline with GPU support
pipeline = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,  # Consider using torch.float32 for GPU computations
    device="cuda",  # Use GPU device if available
).to("cuda")  # Ensure the model is moved to GPU

# Additional pipeline configurations
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config).to("cuda")
pipeline.enable_model_cpu_offload(False)  # Consider not offloading to CPU when using GPU

# Gradio interface setup
demo = gr.Interface(
    fn=transform_image,
    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.",
)

demo.launch()