Create app.py

app.py

@@ -0,0 +1,119 @@
+import streamlit as st
+import torch
+from diffusers import (
+    StableDiffusionXLControlNetPipeline,
+    ControlNetModel
+)
+from PIL import Image
+
+# For demonstration, we’ll assume there’s a published SDXL lineart controlnet on HF.
+# Replace with a valid repo if the name below doesn’t exist or adjust to your needs.
+LINEART_CONTROLNET_REPO = "lllyasviel/sdxl-controlnet-lineart"  # Example placeholder
+SDXL_MODEL_REPO = "RunDiffusion/Juggernaut-XL-v9"               # Or "stabilityai/stable-diffusion-xl-base-1.0"
+
+@st.cache_resource
+def load_pipeline():
+    """
+    Loads the ControlNet model (line-art) and the main Stable Diffusion XL model (Juggernaut XL).
+    Returns a pipeline ready for inference.
+    """
+    controlnet = ControlNetModel.from_pretrained(
+        LINEART_CONTROLNET_REPO,
+        torch_dtype=torch.float16
+    )
+
+    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+        SDXL_MODEL_REPO,
+        controlnet=controlnet,
+        torch_dtype=torch.float16
+    )
+    
+    # Move to GPU if available
+    if torch.cuda.is_available():
+        pipe.to("cuda")
+
+    return pipe
+
+def combine_lineart_and_colormask(lineart: Image.Image, color_mask: Image.Image) -> Image.Image:
+    """
+    Naive example of combining lineart and color mask into a single control image. 
+    Here we just alpha-blend them for demonstration. 
+    In practice, you might want more sophisticated merges, 
+    or you could use multi-ControlNet if the pipelines/models are available.
+    """
+    # Resize color mask to match lineart size
+    color_mask = color_mask.resize(lineart.size)
+    
+    # Convert to RGBA
+    lineart_rgba = lineart.convert("RGBA")
+    color_mask_rgba = color_mask.convert("RGBA")
+    
+    # Simple alpha blend for demonstration
+    blended = Image.blend(lineart_rgba, color_mask_rgba, alpha=0.5)
+    return blended.convert("RGB")
+
+def main():
+    st.title("Line-Art + Color Mask with SDXL ControlNet")
+    st.markdown(
+        "Upload a **line-art sketch** and a **color mask**, then let "
+        "Stable Diffusion XL (Juggernaut XL) + ControlNet (Lineart) do the rest!"
+    )
+
+    # Sidebar inputs for text prompt, etc.
+    prompt = st.sidebar.text_input(
+        "Prompt",
+        value="A cute cartoon-style character with vibrant colors"
+    )
+    negative_prompt = st.sidebar.text_input(
+        "Negative Prompt",
+        value="ugly, deformed"
+    )
+    guidance_scale = st.sidebar.slider(
+        "Guidance Scale (classifier-free)",
+        min_value=1.0,
+        max_value=20.0,
+        value=9.0
+    )
+    num_inference_steps = st.sidebar.slider(
+        "Number of Inference Steps",
+        min_value=10,
+        max_value=100,
+        value=30
+    )
+
+    # Main area for uploading images
+    lineart_file = st.file_uploader("Upload Line-Art Sketch (png/jpg)", type=["png", "jpg", "jpeg"])
+    color_file = st.file_uploader("Upload Color Mask (png/jpg)", type=["png", "jpg", "jpeg"])
+
+    if lineart_file and color_file:
+        lineart_image = Image.open(lineart_file)
+        color_mask = Image.open(color_file)
+
+        st.image(lineart_image, caption="Line-Art Preview", width=300)
+        st.image(color_mask, caption="Color Mask Preview", width=300)
+
+        # Combine images into a single control image
+        combined_control_image = combine_lineart_and_colormask(lineart_image, color_mask)
+        st.image(combined_control_image, caption="Combined Control Image", width=300)
+
+        # Button to run inference
+        if st.button("Generate"):
+            pipe = load_pipeline()
+
+            with st.spinner("Generating image..."):
+                result = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    control_image=combined_control_image,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    # For SDXL pipelines, also pass an additional prompt for the refiner if needed
+                    # refiner_prompt=prompt,  # if your pipeline supports it
+                ).images[0]
+
+            st.image(result, caption="Generated Image", width=512)
+    else:
+        st.warning("Please upload both a line-art sketch and a color mask.")
+
+if __name__ == "__main__":
+    main()