Initialize deployment

Dockerfile

@@ -17,4 +17,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]

requirements.txt

@@ -1,3 +1,16 @@
-altair
-pandas
-streamlit
+opencv-python-headless
+numpy
+tqdm
+torch
+torchvision
+diffusers
+transformers
+accelerate
+controlnet-aux
+datasets
+torch-fidelity
+lpips
+pillow
+hf_xet
+streamlit
+huggingface-hub

src/config.py

@@ -0,0 +1,9 @@
+import os
+import torch
+from huggingface_hub import hf_hub_download
+
+app_controlnet_path = "SaitoHoujou/Finetuned-ControlNet_for_Diffusion-Model"
+app_base_model = "botp/stable-diffusion-v1-5" 
+app_hed_model = 'lllyasviel/Annotators'
+app_device = "cuda" if torch.cuda.is_available() else "cpu"
+app_dtype = torch.float16 if device == "cuda" else torch.float32

src/inference.py

@@ -0,0 +1,34 @@
+import torch
+from PIL import Image
+
+def generate_image(
+    pipe, 
+    hed, 
+    input_image: Image.Image, 
+    prompt: str, 
+    neg_prompt: str, 
+    guidance_scale: float, 
+    control_scale: float, 
+    device: str,
+    seed: int = 42
+) -> (Image.Image, Image.Image):
+    
+    condition_image = hed(
+        input_image, 
+        detect_resolution=512, 
+        image_resolution=512
+    )
+    
+    generator = torch.Generator(device=device).manual_seed(seed)
+    
+    output_image = pipe(
+        prompt=prompt,
+        negative_prompt=neg_prompt,
+        image=condition_image,
+        num_inference_steps=30,
+        generator=generator,
+        guidance_scale=guidance_scale,
+        controlnet_conditioning_scale=control_scale
+    ).images[0]
+    
+    return output_image, condition_image

src/loader.py

@@ -0,0 +1,29 @@
+import streamlit as st
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from controlnet_aux import HEDdetector
+import config 
+
+@st.cache_resource
+def load_hed_detector():
+    hed = HEDdetector.from_pretrained(config.app_hed_model)
+    hed = hed.to(config.app_device)
+    return hed
+
+@st.cache_resource
+def load_pipeline():
+    controlnet = ControlNetModel.from_pretrained(
+        config.app_controlnet_path, 
+        torch_dtype=config.app_dtype
+    )
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        config.app_base_model, 
+        controlnet=controlnet, 
+        torch_dtype=config.app_dtype,
+        safety_checker=None
+    )
+    
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to(config.app_device)
+    return pipe

src/streamlit_app.py

@@ -1,40 +0,0 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))

streamlit_app.py

@@ -0,0 +1,66 @@
+import streamlit as st
+from PIL import Image
+from src import config, loader, inference
+
+st.set_page_config(layout="wide")
+st.title("Demo Generating Image from Face Sketch ")
+
+st.markdown("---")
+
+with st.spinner("Loading models..."):
+    pipe = loader.load_pipeline()
+    hed = loader.load_hed_detector()
+
+with st.sidebar:
+    st.header("Configuration")
+    prompt = st.text_area(
+        "Positive Prompt",
+        "(hyper-realistic photo:1.2), (ultra-detailed skin texture:1.1), detailed pores, realistic eyes, (Caucasian man:1.1), sharp focus, 8k UHD, professional studio lighting",
+        height=100
+    )
+    negative_prompt = st.text_area(
+        "Negative Prompt",
+        "(drawing:1.4), (sketch:1.4), (painting:1.3), cartoon, 3D, render, CGI, anime, illustration, (deformed:1.2), (disfigured:1.2), ugly, bad anatomy, (blurry:1.1)",
+        height=100
+    )
+    guidance_scale = st.slider("Guidance Scale", 1.0, 15.0, 8.0, 0.5)
+    control_scale = st.slider("ControlNet Scale", 0.0, 1.0, 0.9, 0.1)
+    uploaded_file = st.file_uploader("Upload your face sketch here...", type=["png", "jpg", "jpeg"])
+    run_button = st.button("Generate Image", width='stretch', type="primary")
+    
+col1, col2 = st.columns(2)
+
+with col1:
+    st.header("1. Uploaded Face Sketch")
+    st.markdown("Input Face Sketch")
+    input_image_placeholder = st.empty()
+    st.markdown("HED Image")
+    hed_image_placeholder = st.empty()
+
+with col2:
+    st.header("2. Generated Image")
+    
+    st.markdown("Output Generated Image")
+    output_image_placeholder = st.empty()
+
+    if run_button and uploaded_file:
+        input_image = Image.open(uploaded_file).convert("RGB")
+        input_image_placeholder.image(input_image, caption="Uploaded sketch", width='stretch')
+
+        with st.spinner("Generating image..."):
+            output_image, condition_image = inference.generate_image(
+                pipe=pipe,
+                hed=hed,
+                input_image=input_image,
+                prompt=prompt,
+                neg_prompt=negative_prompt,
+                guidance_scale=guidance_scale,
+                control_scale=control_scale,
+                device=config.app_device
+            )
+
+        hed_image_placeholder.image(condition_image, caption="HED Image", width='stretch')
+        output_image_placeholder.image(output_image, caption="Generated Image", width='stretch')
+
+    elif uploaded_file:
+        input_image_placeholder.info("Click the generate image button on sidebar")