Spaces:
Running
Running
| import os | |
| os.environ["HF_HOME"] = "/data/huggingface" | |
| os.environ["TRANSFORMERS_CACHE"] = "/data/huggingface" | |
| os.makedirs("/data/huggingface/hub", exist_ok=True) | |
| # os.makedirs("/data/huggingface/clip_vision_model", exist_ok=True) | |
| import torch | |
| from diffusers import StableDiffusionImg2ImgPipeline | |
| from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection | |
| from PIL import Image | |
| # --- Place any download or path setup here --- old | |
| # MODEL_ID ="runwayml/stable-diffusion-v1-5" # Can swap for custom path if using IP-Adapter | |
| # ADAPTER_PATH = "/workspace/.cache/huggingface/ip_adapter/ip-adapter_sd15.bin" | |
| # ADAPTER_DIR = "/workspace/.cache/huggingface/ip_adapter" | |
| # DEVICE = "cpu" | |
| # MODEL_CACHE = "/workspace/.cache/huggingface" | |
| # ---- SETTINGS ---- | |
| MODEL_ID = "runwayml/stable-diffusion-v1-5" | |
| IPADAPTER_REPO = "h94/IP-Adapter" | |
| IPADAPTER_WEIGHT_NAME = "ip-adapter_sd15.bin" | |
| DEVICE = "cpu" # Change to "cuda" if you have GPU | |
| CACHE_DIR = os.environ.get("HF_HOME", "/data/huggingface") | |
| # (Optional) Download IP-Adapter weights and patch pipeline if desired | |
| # Load the model ONCE at startup, not per request! | |
| pipe = StableDiffusionImg2ImgPipeline.from_pretrained( | |
| MODEL_ID, | |
| torch_dtype=torch.float32, | |
| cache_dir=CACHE_DIR, | |
| # safety_checker=None, # Disable for demo/testing; enable in prod | |
| ).to(DEVICE) | |
| pipe.load_ip_adapter( | |
| pretrained_model_name_or_path_or_dict=IPADAPTER_REPO, | |
| subfolder="models", | |
| weight_name=IPADAPTER_WEIGHT_NAME | |
| ) | |
| # # Load vision encoder and processor for IP-Adapter embedding | |
| # vision_encoder = CLIPVisionModelWithProjection.from_pretrained( | |
| # "h94/IP-Adapter", # repo_id (main IP-Adapter repo) | |
| # subfolder="clip_vision_model",# subfolder within the repo! | |
| # cache_dir=CACHE_DIR | |
| # ) | |
| # image_processor = CLIPImageProcessor.from_pretrained( | |
| # "h94/IP-Adapter", | |
| # subfolder="clip_vision_model", | |
| # cache_dir=CACHE_DIR | |
| # ) | |
| def generate_sticker(input_image: Image.Image, style: str ): | |
| """ | |
| 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) | |
| face_img = input_image.convert("RGB").resize((512, 512)) | |
| # inputs = image_processor(images=face_img, return_tensors="pt").to(DEVICE) | |
| # with torch.no_grad(): | |
| # image_embeds = vision_encoder(**inputs).image_embeds | |
| # 2. Prepare image for SD pipeline | |
| init_image = input_image.convert("RGB").resize((512, 512)) | |
| # IP-Adapter expects the reference image via image_embeds, which is produced by this function: | |
| # image_embeds = pipe.prepare_ip_adapter_image_embeds(face_img) | |
| prompt = (f"A set of twelve {style}-style digital stickers" | |
| "each with a different expression: laughing, angry, crying, sulking, thinking, sleepy, blowing a kiss, winking, surprised, happy, sad, and confused. " | |
| "Each sticker has a bold black outline and a transparent background, in a playful, close-up cartoon style." | |
| ) | |
| # Run inference (low strength for identity preservation) | |
| result = pipe( | |
| prompt=prompt, | |
| image=init_image, | |
| # image_embeds=image_embeds, | |
| ip_adapter_image=face_img, | |
| strength=0.6, | |
| guidance_scale=8, | |
| num_inference_steps=40 | |
| ) | |
| # Return the generated image (as PIL) | |
| return result.images[0] | |