import gradio as gr
from diffusers import StableDiffusionImg2ImgPipeline
import torch
from PIL import Image
import io
import logging
import numpy as np

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load the Stable Diffusion image-to-image pipeline
model_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe = pipe.to("cpu")  # Use CPU (slow, switch to "cuda" if available)

def transform_profile_images(files):
    try:
        if not files:
            logger.info("No files uploaded.")
            return []

        # Prompt for LinkedIn-style professional images
        prompt = ("A professional LinkedIn profile picture, clear background, business attire, "
                  "high resolution, photorealistic, studio lighting")

        transformed_images = []
        for i, file in enumerate(files):
            logger.info("Processing file %d", i + 1)
            try:
                image = Image.open(io.BytesIO(file))
                logger.info(f"Original image type: {type(image)}")
            except Exception as e:
                logger.error(f"Failed to open image: {e}")
                continue

            if image.mode != "RGB":
                image = image.convert("RGB")
                logger.info("Converted image to RGB mode.")

            if not isinstance(image, Image.Image):
                logger.error("Image is not a PIL.Image.Image object.")
                continue

            image = image.resize((512, 512))  # Better resolution for processing
            logger.info("Resized image to 512x512.")

            try:
                output = pipe(prompt=prompt, image=image, strength=0.7, guidance_scale=7.5).images[0]
                transformed_images.append(output)
                logger.info("Successfully processed file %d", i + 1)
            except Exception as e:
                logger.error(f"Failed to process image: {e}")

        return transformed_images
    except Exception as e:
        logger.error(f"An error occurred: {e}")
        return []

# Gradio Interface
demo = gr.Interface(
    fn=transform_profile_images,
    inputs=gr.File(label="Upload Your Profile Images", file_count="multiple", type="binary"),
    outputs=gr.Gallery(label="Professional LinkedIn Profile Images", columns=2, height="auto"),
    title="Professional LinkedIn Profile Picture Generator",
    description=(
        "Upload one or more casual profile images and let our AI transform them into professional-looking "
        "LinkedIn portraits. The system uses a Stable Diffusion image-to-image model to refine your images "
        "with business attire, a clear background, and studio lighting."
    )
)

if __name__ == "__main__":
    demo.launch(share=True)