give ability to have multiple models

app.py

@@ -3,6 +3,7 @@ import numpy as np
 import os
 import spaces
 import random
+import json
 # from image_gen_aux import DepthPreprocessor
 from PIL import Image
 import torch
@@ -16,73 +17,34 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
 pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
-pipe.load_lora_weights("Himanshu806/testLora")
-pipe.enable_lora()
-
-# vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae")
-# processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
-
-# preprocess = transforms.Compose(
-#     [
-#         transforms.Resize(
-#             (vae.config.sample_size, vae.config.sample_size),
-#             interpolation=transforms.InterpolationMode.BILINEAR,
-#         ),
-#         transforms.ToTensor(),
-#         transforms.Normalize([0.5], [0.5]),
-#     ]
-# )
-#
-# image_np = image[0].cpu().numpy()  # Move to CPU and convert to NumPy
-
-# if image_np.shape[0] == 3:  # Check if channels are first
-#     image_np = image_np.transpose(1, 2, 0)
-
-# image_np = (image_np * 255).astype(np.uint8)
-
-# image = Image.fromarray(image_np)
-
-# def calculate_optimal_dimensions(image: Image.Image):
-#     # Extract the original dimensions
-#     original_width, original_height = image.size
-
-#     # Set constants
-#     MIN_ASPECT_RATIO = 9 / 16
-#     MAX_ASPECT_RATIO = 16 / 9
-#     FIXED_DIMENSION = 1024
-
-#     # Calculate the aspect ratio of the original image
-#     original_aspect_ratio = original_width / original_height
-
-#     # Determine which dimension to fix
-#     if original_aspect_ratio > 1:  # Wider than tall
-#         width = FIXED_DIMENSION
-#         height = round(FIXED_DIMENSION / original_aspect_ratio)
-#     else:  # Taller than wide
-#         height = FIXED_DIMENSION
-#         width = round(FIXED_DIMENSION * original_aspect_ratio)
-
-#     # Ensure dimensions are multiples of 8
-#     width = (width // 8) * 8
-#     height = (height // 8) * 8
-
-#     # Enforce aspect ratio limits
-#     calculated_aspect_ratio = width / height
-#     if calculated_aspect_ratio > MAX_ASPECT_RATIO:
-#         width = (height * MAX_ASPECT_RATIO // 8) * 8
-#     elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
-#         height = (width / MIN_ASPECT_RATIO // 8) * 8
-
-#     # Ensure width and height remain above the minimum dimensions
-#     width = max(width, 576) if width == FIXED_DIMENSION else width
-#     height = max(height, 576) if height == FIXED_DIMENSION else height
-
-#     return width, height
+# pipe.load_lora_weights("Himanshu806/testLora")
+# pipe.enable_lora()
+
+with open("lora_models.json", "r") as f:
+    lora_models = json.load(f)
+
+def download_model(model_name, model_path):
+    print(f"Downloading model: {model_name} from {model_path}")
+    try:
+        pipe.load_lora_weights(model_path)
+        print(f"Successfully downloaded model: {model_name}")
+    except Exception as e:
+        print(f"Failed to download model: {model_name}. Error: {e}")
+
+# Iterate through the models and download each one
+for model_name, model_path in lora_models.items():
+    download_model(model_name, model_path)
+
+lora_models["None"] = None
 
 @spaces.GPU(durations=300)
-def infer(edit_images, prompt, prompt2, width, height, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(edit_images, prompt, width, height, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     # pipe.enable_xformers_memory_efficient_attention()
 
+    if lora_model != "None":
+        pipe.load_lora_weights(lora_models[lora_model])
+        pipe.enable_lora()
+
     image = edit_images["background"]
     # width, height = calculate_optimal_dimensions(image)
     mask = edit_images["layers"][0]
@@ -93,7 +55,7 @@ def infer(edit_images, prompt, prompt2, width, height, seed=42, randomize_seed=F
     image = pipe(
         # mask_image_latent=vae.encode(controlImage),
         prompt=prompt,
-        prompt_2=prompt2,
+        prompt_2=prompt,
         image=image,
         mask_image=mask,
         height=height,
@@ -147,13 +109,13 @@ with gr.Blocks(css=css) as demo:
                     placeholder="Enter your prompt",
                     container=False,
                 )
-                prompt2 = gr.Text(
-                    label="Prompt2",
-                    show_label=False,
-                    max_lines=2,
-                    placeholder="Enter your second prompt",
-                    container=False,
+
+                lora_model = gr.Dropdown(
+                    label="Select LoRA Model",
+                    choices=list(lora_models.keys()),
+                    value="None",
                 )
+
                 run_button = gr.Button("Run")
 
             result = gr.Image(label="Result", show_label=False)
@@ -209,7 +171,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [edit_image, prompt,prompt2, width, height, seed, randomize_seed, guidance_scale, num_inference_steps],
+        inputs = [edit_image, prompt, width, height, lora_model, seed, randomize_seed, guidance_scale, num_inference_steps],
         outputs = [result, seed]
     )
 

lora_models.json

@@ -0,0 +1,4 @@
+{
+  "RahulFineTuned": "Himanshu806/testLora",
+  "KodaRealistic": "alvdansen/flux-koda"
+}