Update worker_runpod.py

worker_runpod.py

@@ -23,7 +23,7 @@ from huggingface_hub import HfApi, HfFolder
 tokenxf = os.getenv("HF_API_TOKEN")
 # Low GPU memory mode
 low_gpu_memory_mode = False
-
+lora_path = "/content/shirtlift.safetensors"
 def download_image(url, download_dir="asset"):
     # Ensure the download directory exists
     if not os.path.exists(download_dir):
@@ -121,6 +121,9 @@ def generate(input):
     validation_image_end = values.get("validation_image_end", None)
     
     generator = torch.Generator(device="cuda").manual_seed(seed)
+    if lora_path is not None:
+        pipeline = merge_lora(pipeline, lora_path, lora_weight)
+    
     aspect_ratio_sample_size = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
     start_img = Image.open(downloaded_image_path)
     original_width, original_height = start_img.size