app.py

import torch
from diffusers import AutoPipelineForImage2Image
from PIL import Image
import gradio as gr

# Initialize the image-to-image pipeline
pipeline = AutoPipelineForImage2Image.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)

# Enable CPU offloading to manage memory usage
pipeline.enable_model_cpu_offload()

# Function to process the image and prompt, and generate a new image
def process_image(prompt, init_image):
    # Convert the uploaded file to an image
    init_image = Image.open(init_image).convert("RGB")

    # Generate an image based on the prompt and initial image
    with torch.no_grad():  # This ensures that no gradients are calculated, saving memory
        generated_image = pipeline(prompt, init_image=init_image, num_inference_steps=50).images[0]

    return generated_image

# Create the Gradio interface
interface = gr.Interface(
    fn=process_image,
    inputs=[gr.inputs.Textbox(label="Prompt"), gr.inputs.Image(label="Initial Image", type="pil")],
    outputs=gr.outputs.Image(type="pil", label="Generated Image"),
    title="Image-to-Image Transformation",
    description="Upload an image and enter a prompt to transform the image accordingly."
)

# Launch the interface
interface.launch()