Update app.py

app.py

@@ -1,28 +1,38 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline
-import os
-from huggingface_hub import login
-from safetensors.torch import load_file
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+from diffusers.utils import load_image
 
-# Login using the token
-login(token=os.environ.get("HF_TOKEN"))
-
-# Initialize the pipeline with CPU
+# Initialize the base pipeline
 model_id = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionPipeline.from_pretrained(
     model_id,
-    torch_dtype=torch.float32,
-    use_auth_token=os.environ.get("HF_TOKEN")
+    torch_dtype=torch.float32
+)
+
+# Load the ControlNet UNet
+unet = UNet2DConditionModel.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", 
+    subfolder="unet",
+    torch_dtype=torch.float32
 )
+unet.add_extra_conditions(["canny"])
 
-# Load your LoRA weights from local file
-lora_path = "models/40kHalf.safetensors"  # Update this path to where you upload the file in your Space
-state_dict = load_file(lora_path)
-pipe.unet.load_attn_procs(state_dict)
+# Add the conditioning
+pipe.unet = unet
+
+# Load the ControlLoRA weights
+pipe.load_lora_weights(
+    "models",  # Local path to your ControlLoRA weights
+    weight_name="40kHalf.safetensors",
+    adapter_name="control_lora"
+)
 
 def generate_image(prompt, negative_prompt, guidance_scale, steps):
     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,