update code

Dockerfile

@@ -9,7 +9,6 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 COPY requirements.txt ./
-COPY src/ ./src/
 COPY streamlit_app.py ./
 
 RUN pip3 install -r requirements.txt

src/config.py

@@ -1,9 +0,0 @@
-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 app_device == "cuda" else torch.float32

src/inference.py

@@ -1,34 +0,0 @@
-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

@@ -1,29 +0,0 @@
-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

streamlit_app.py

@@ -1,6 +1,70 @@
 import streamlit as st
 from PIL import Image
-from src import config, loader, inference
+import torch
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
+from controlnet_aux import HEDdetector
+
+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 app_device == "cuda" else torch.float32
+
+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
+
+@st.cache_resource
+def load_hed_detector():
+    hed = HEDdetector.from_pretrained(app_hed_model)
+    hed = hed.to(app_device)
+    return hed
+
+@st.cache_resource
+def load_pipeline():
+    controlnet = ControlNetModel.from_pretrained(
+        app_controlnet_path, 
+        torch_dtype=app_dtype
+    )
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        app_base_model, 
+        controlnet=controlnet, 
+        torch_dtype=app_dtype,
+        safety_checker=None
+    )
+    
+    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to(app_device)
+    return pipe
 
 st.set_page_config(layout="wide")
 st.title("Demo Generating Image from Face Sketch ")
@@ -8,8 +72,8 @@ st.title("Demo Generating Image from Face Sketch ")
 st.markdown("---")
 
 with st.spinner("Loading models..."):
-    pipe = loader.load_pipeline()
-    hed = loader.load_hed_detector()
+    pipe = load_pipeline()
+    hed = load_hed_detector()
 
 with st.sidebar:
     st.header("Configuration")
@@ -48,7 +112,7 @@ with col2:
         input_image_placeholder.image(input_image, caption="Uploaded sketch", width='stretch')
 
         with st.spinner("Generating image..."):
-            output_image, condition_image = inference.generate_image(
+            output_image, condition_image = generate_image(
                 pipe=pipe,
                 hed=hed,
                 input_image=input_image,
@@ -56,7 +120,7 @@ with col2:
                 neg_prompt=negative_prompt,
                 guidance_scale=guidance_scale,
                 control_scale=control_scale,
-                device=config.app_device
+                device=app_device
             )
 
         hed_image_placeholder.image(condition_image, caption="HED Image", width='stretch')