Create app.py

app.py

@@ -0,0 +1,147 @@
+# app.py
+import os
+import io
+import random
+from PIL import Image, ImageOps
+import numpy as np
+import streamlit as st
+
+# ---- ML libs ----
+import torch
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, UniPCMultistepScheduler
+from huggingface_hub import login
+
+# ---- OpenCV for simple preproc (Canny) ----
+import cv2
+
+st.set_page_config(page_title="Sketch2Face (Streamlit + ControlNet)", layout="centered")
+
+st.title("Sketch2Face — turn your face sketches into stylized images")
+st.write("Upload a face sketch (line drawing). Use the prompt to guide style, pose & mood.")
+
+# Get HF token (recommended to set as secret on HF Spaces or as env var locally)
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+if HF_TOKEN:
+    try:
+        login(token=HF_TOKEN)
+    except Exception:
+        pass
+else:
+    st.warning("No Hugging Face token found. On Spaces, add HF_TOKEN in Settings → Secrets for model download. Locally use 'huggingface-cli login'.")
+
+# Sidebar controls
+with st.sidebar:
+    st.header("Generation settings")
+    model_id = st.text_input("Stable Diffusion model (hf repo)", value="runwayml/stable-diffusion-v1-5")
+    controlnet_id = st.text_input("ControlNet (canny) repo", value="lllyasviel/sd-controlnet-canny")
+    prompt = st.text_area("Prompt", value="A realistic portrait of a young man, soft lighting, cinematic")
+    negative_prompt = st.text_area("Negative prompt (optional)", value="lowres, deformed, extra fingers, watermark")
+    guidance_scale = st.slider("Guidance scale", 1.0, 20.0, 7.5)
+    strength = st.slider("Strength (how much to change sketch)", 0.1, 1.0, 0.7)
+    num_inference_steps = st.slider("Steps", 10, 60, 28)
+    seed = st.number_input("Seed (0 for random)", min_value=0, max_value=999999999, value=0, step=1)
+    use_gpu = st.checkbox("Use GPU (if available)", value=True)
+    run_btn = st.button("Generate")
+
+# Upload sketch
+uploaded = st.file_uploader("Upload your sketch (png/jpg). Prefer simple line art.", type=["png","jpg","jpeg"])
+example_col1, example_col2 = st.columns(2)
+with example_col1:
+    st.markdown("**Tip:**** clear black lines on white background work best.")
+with example_col2:
+    st.markdown("**Tip:** crop to face / 1:1 or 3:4 ratio.")
+
+@st.cache_resource(show_spinner=False)
+def load_models(sd_model_id, cn_model_id, device):
+    # Load ControlNet then the combined pipeline
+    controlnet = ControlNetModel.from_pretrained(
+        cn_model_id, torch_dtype=torch.float16 if device=="cuda" else torch.float32
+    )
+    pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+        sd_model_id,
+        controlnet=controlnet,
+        safety_checker=None,
+        torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+    )
+    # Scheduler & device
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    if device == "cuda":
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.to("cuda")
+    else:
+        pipe.to("cpu")
+    return pipe
+
+def prepare_control_image_pil(pil_img, target_size=512):
+    # Ensure grayscale -> convert to single-channel edge map using Canny
+    img = pil_img.convert("RGB")
+    open_cv_image = np.array(img)[:, :, ::-1]  # RGB->BGR
+    gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)
+    # Auto-threshold can be useful; here we use fixed, but you can expose sliders
+    edges = cv2.Canny(gray, 100, 200)
+    edges = cv2.resize(edges, (target_size, target_size))
+    # convert single channel to 3-channel PIL
+    edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
+    return Image.fromarray(edges_rgb)
+
+def prepare_init_image(pil_img, target_size=512):
+    img = pil_img.convert("RGB")
+    img = ImageOps.fit(img, (target_size, target_size), Image.LANCZOS)
+    return img
+
+if run_btn:
+    if not uploaded:
+        st.error("Please upload a sketch first.")
+    else:
+        device = "cuda" if (torch.cuda.is_available() and use_gpu) else "cpu"
+        with st.spinner("Loading models (first run may take ~1-2 minutes)..."):
+            pipe = load_models(model_id, controlnet_id, device)
+        # load user image
+        img = Image.open(uploaded)
+        control_image = prepare_control_image_pil(img, target_size=512)
+        init_image = prepare_init_image(img, target_size=512)
+
+        # seed
+        gen_seed = None if seed == 0 else int(seed)
+        generator = torch.Generator(device=device)
+        if gen_seed is not None:
+            generator = generator.manual_seed(gen_seed)
+        else:
+            generator = None
+
+        with st.spinner("Generating..."):
+            try:
+                output = pipe.img2img(
+                    prompt=prompt,
+                    image=init_image,
+                    control_image=control_image,
+                    negative_prompt=negative_prompt or None,
+                    strength=float(strength),
+                    guidance_scale=float(guidance_scale),
+                    num_inference_steps=int(num_inference_steps),
+                    generator=generator,
+                )
+            except Exception as e:
+                st.exception(f"Generation failed: {e}")
+                raise
+
+        result = output.images[0]
+        st.image(result, caption="Generated image", use_column_width=True)
+        # offer download
+        buf = io.BytesIO()
+        result.save(buf, format="PNG")
+        buf.seek(0)
+        st.download_button("Download image (PNG)", data=buf, file_name="sketch2face.png", mime="image/png")
+
+# Show sample control image / debug
+if uploaded:
+    try:
+        img = Image.open(uploaded)
+        control_img = prepare_control_image_pil(img, target_size=256)
+        st.caption("Preview: internal Canny/control image (what ControlNet sees)")
+        st.image(control_img)
+    except Exception:
+        pass
+
+st.markdown("---")
+st.markdown("Made for sketch-to-face. Adjust prompt & strength. For best results, upload clear line sketches and try style prompts like 'photorealistic', 'studio lighting', or artists' names (check licenses).")