import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

# --- Place any download or path setup here ---
MODEL_ID = "runwayml/stable-diffusion-v1-5"  # Can swap for custom path if using IP-Adapter
DEVICE = "cpu"
MODEL_CACHE = "./models"

# (Optional) Download IP-Adapter weights and patch pipeline if desired

def generate_sticker(input_image, prompt):
    """
    Given a user image and a prompt, generates a sticker/emoji-style portrait.
    """
    # Load the model (download if not present)
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        cache_dir=MODEL_CACHE,
        safety_checker=None,  # Disable for demo/testing
    ).to(DEVICE)

    # Preprocess the image (resize, etc)
    init_image = input_image.convert("RGB").resize((512, 512))

    # Run inference (low strength for identity preservation)
    result = pipe(
        prompt=prompt,
        image=init_image,
        strength=0.65,
        guidance_scale=7.5,
        num_inference_steps=30
    )
    # Return the generated image (as PIL)
    return result.images[0]