Sticker_Diffusion / utils.py
KarthikAI's picture
Upload 5 files
d646931 verified
raw
history blame
1.14 kB
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]