Spaces:
Running
Running
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] | |