Update controlnet/callable_functions.py

controlnet/callable_functions.py

@@ -10,7 +10,7 @@ from transformers import AutoProcessor, SiglipVisionModel
 
 
 
-def process_single_image(image_path, image=None):
+def process_single_image(model,image_path, image=None):
     
     # Set up model components
     unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", torch_dtype=torch.float16, device="cuda")
@@ -19,7 +19,7 @@ def process_single_image(image_path, image=None):
     stylecodes_model= stylecodes_model.to("cuda")   
 
 
-    stylecodes_model.load_model("models/stylecodes_sd15_v1.bin")
+    stylecodes_model.load_model(model)
     # Load and preprocess image
     if image is None:
         image = Image.open(image_path).convert("RGB")
@@ -40,7 +40,7 @@ def process_single_image(image_path, image=None):
     return code
 
 
-def process_single_image_both_ways(image_path, prompt, num_inference_steps,image=None):
+def process_single_image_both_ways(model,image_path, prompt, num_inference_steps,image=None):
     # Load and preprocess image
     # Set up model components
     unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", torch_dtype=torch.float16, device="cuda")
@@ -56,7 +56,7 @@ def process_single_image_both_ways(image_path, prompt, num_inference_steps,image
         steps_offset=1,
     )
 
-    stylecodes_model.load_model("models/stylecodes_sd15_v1.bin")
+    stylecodes_model.load_model(model)
 
     pipe = StableDiffusionPipelineXSv2.from_pretrained(
         "runwayml/stable-diffusion-v1-5",
@@ -95,7 +95,7 @@ def process_single_image_both_ways(image_path, prompt, num_inference_steps,image
     # Save the output image
 
 
-def make_stylecode(image_path, image=None):
+def make_stylecode(model,image_path, image=None):
     
     # Set up model components
     unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet", torch_dtype=torch.float16, device="cuda")
@@ -104,7 +104,7 @@ def make_stylecode(image_path, image=None):
     stylecodes_model= stylecodes_model.to("cuda")   
 
 
-    stylecodes_model.load_model("models/stylecodes_sd15_v1.bin")
+    stylecodes_model.load_model(model)
     # Load and preprocess image
     if image is None:
         image = Image.open(image_path).convert("RGB")