| |
| import os |
| import streamlit as st |
| from PIL import Image |
| import torch |
| from diffusers import PaintByExamplePipeline |
|
|
| |
| os.environ["HF_HOME"] = "/app/.cache/huggingface" |
| os.environ["TRANSFORMERS_CACHE"] = "/app/.cache/transformers" |
| os.environ["DIFFUSERS_CACHE"] = "/app/.cache/diffusers" |
|
|
| os.makedirs("/app/.cache/huggingface", exist_ok=True) |
| os.makedirs("/app/.cache/transformers", exist_ok=True) |
| os.makedirs("/app/.cache/diffusers", exist_ok=True) |
|
|
| st.set_page_config(page_title="Paint-By-Example", layout="centered") |
|
|
| @st.cache_resource(show_spinner=False) |
| def load_pipe(): |
| dtype = torch.float16 if torch.cuda.is_available() else torch.float32 |
| pipe = PaintByExamplePipeline.from_pretrained( |
| "Fantasy-Studio/Paint-by-Example", |
| torch_dtype=dtype, |
| cache_dir=os.environ["DIFFUSERS_CACHE"] |
| ) |
| pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu") |
| return pipe |
|
|
| pipe = load_pipe() |
|
|
| st.title("🧩 Paint-By-Example (édition par image de référence)") |
| img_file = st.file_uploader("Image d’entrée", type=["png","jpg","jpeg"]) |
| mask_file = st.file_uploader("Masque (blanc = à modifier)", type=["png","jpg","jpeg"]) |
| ref_file = st.file_uploader("Image de référence (style/objet à copier)", type=["png","jpg","jpeg"]) |
| steps = st.slider("Steps", 1, 100, 50) |
|
|
| if st.button("Appliquer"): |
| if not (img_file and mask_file and ref_file): |
| st.warning("Charge l’image, le masque et l’image de référence.") |
| else: |
| image = Image.open(img_file).convert("RGB") |
| mask = Image.open(mask_file).convert("L") |
| ref = Image.open(ref_file).convert("RGB") |
|
|
| with st.spinner("Génération…"): |
| out = pipe( |
| image=image, |
| mask_image=mask, |
| example_image=ref, |
| num_inference_steps=steps |
| ).images[0] |
| st.image(out, caption="Résultat", use_column_width=True) |
|
|
