Update app.py

app.py

@@ -1,63 +1,98 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, UNet2DConditionModel
-from diffusers.utils import load_image
-
-# Initialize the base pipeline
-model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
-    torch_dtype=torch.float32
-)
+import numpy as np
+import cv2
+from diffusers import StableDiffusionPipeline
+from model import UNet2DConditionModelEx, StableDiffusionControlLoraV3Pipeline
+from PIL import Image
+import os
+from huggingface_hub import login
+
+# Login using the token
+login(token=os.environ.get("HF_TOKEN"))
+
+# Initialize the models
+base_model = "runwayml/stable-diffusion-v1-5"
+dtype = torch.float32
 
-# Load the ControlNet UNet
-unet = UNet2DConditionModel.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", 
+# Load the custom UNet
+unet = UNet2DConditionModelEx.from_pretrained(
+    base_model,
     subfolder="unet",
-    torch_dtype=torch.float32
+    torch_dtype=dtype
 )
-unet.add_extra_conditions(["canny"])
 
-# Add the conditioning
-pipe.unet = unet
+# Add conditioning
+unet = unet.add_extra_conditions("ow-gbi-control-lora")
+
+# Create the pipeline with custom UNet
+pipe = StableDiffusionControlLoraV3Pipeline.from_pretrained(
+    base_model, 
+    unet=unet,
+    torch_dtype=dtype
+)
 
 # Load the ControlLoRA weights
 pipe.load_lora_weights(
-    "models",  # Local path to your ControlLoRA weights
-    weight_name="40kHalf.safetensors",
-    adapter_name="control_lora"
+    "models",
+    weight_name="40kHalf.safetensors"
 )
 
-def generate_image(prompt, negative_prompt, guidance_scale, steps):
+def get_canny_image(image, low_threshold=100, high_threshold=200):
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
+    if image.shape[2] == 4:
+        image = image[..., :3]
+    
+    canny_image = cv2.Canny(image, low_threshold, high_threshold)
+    canny_image = np.stack([canny_image] * 3, axis=-1)
+    return Image.fromarray(canny_image)
+
+def generate_image(input_image, prompt, negative_prompt, guidance_scale, steps, low_threshold, high_threshold):
+    canny_image = get_canny_image(input_image, low_threshold, high_threshold)
+    
     with torch.no_grad():
-        # Get canny edges (you'll need to implement this)
-        # For now, let's assume the input image processing is handled separately
-        
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=steps,
             guidance_scale=guidance_scale,
+            image=canny_image
         ).images[0]
-    return image
+    
+    return canny_image, image
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
+            input_image = gr.Image(label="Input Image", type="numpy")
             prompt = gr.Textbox(label="Prompt")
             negative_prompt = gr.Textbox(label="Negative Prompt")
+            with gr.Row():
+                low_threshold = gr.Slider(minimum=1, maximum=255, value=100, label="Canny Low Threshold")
+                high_threshold = gr.Slider(minimum=1, maximum=255, value=200, label="Canny High Threshold")
             guidance_scale = gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale")
             steps = gr.Slider(minimum=1, maximum=100, value=50, label="Steps")
             generate = gr.Button("Generate")
         
         with gr.Column():
+            canny_output = gr.Image(label="Canny Edge Detection")
             result = gr.Image(label="Generated Image")
     
     generate.click(
         fn=generate_image,
-        inputs=[prompt, negative_prompt, guidance_scale, steps],
-        outputs=result
+        inputs=[
+            input_image,
+            prompt,
+            negative_prompt,
+            guidance_scale,
+            steps,
+            low_threshold,
+            high_threshold
+        ],
+        outputs=[canny_output, result]
     )
 
 demo.launch()