Update app.py

app.py

@@ -1,4 +1,3 @@
-import spaces
 import gradio as gr
 import torch
 from PIL import Image
@@ -10,7 +9,7 @@ import numpy as np
 
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
+    img.save(unique_index)
     return unique_name
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
@@ -26,11 +25,16 @@ if not torch.cuda.is_available():
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 
-lora_repo = "strangerzonehf/Flux-Realmix2"
-trigger_word = "Realmix v2"  # Leave trigger_word blank if not used.
-
-pipe.load_lora_weights(lora_repo)
-pipe.to("cuda")
+LORA_OPTIONS = {
+    "Sketch-Smudge": {
+        "lora_repo": "strangerzonehf/Flux-Sketch-Smudge-LoRA",
+        "trigger_word": "Sketch Smudge"
+    },
+    "Sketch-Sized": {
+        "lora_repo": "strangerzonehf/Flux-Sketch-Sized-LoRA",
+        "trigger_word": "Sketch Sized"
+    }
+}
 
 style_list = [
     {
@@ -68,12 +72,20 @@ def generate(
     guidance_scale: float = 3,
     randomize_seed: bool = False,
     style_name: str = DEFAULT_STYLE_NAME,
+    lora_choice: str = "Sketch-Smudge",
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
 
     positive_prompt = apply_style(style_name, prompt)
     
+    # Load the selected LoRA weights
+    lora_repo = LORA_OPTIONS[lora_choice]["lora_repo"]
+    trigger_word = LORA_OPTIONS[lora_choice]["trigger_word"]
+    
+    pipe.load_lora_weights(lora_repo)
+    pipe.to("cuda")
+
     if trigger_word:
         positive_prompt = f"{trigger_word} {positive_prompt}"
 
@@ -91,7 +103,6 @@ def generate(
     return image_paths, seed
 
 examples = [
-
     "midjourney mix, a tiny astronaut hatching from an egg on the moon",
     "midjourney mix, intense Red, a black cat is facing the left side of the frame. The cats head is tilted upward, with its eyes closed. Its whiskers are protruding from its mouth, adding a touch of warmth to the scene. The background is a vibrant red, creating a striking contrast with the cats fur.",
     "midjourney mix, a close-up shot of a womans face, the womans hair is wet, and she is wearing a cream-colored sweater. The background is blurred, and there are red and white signs visible in the background. The womans eyebrows are wet, adding a touch of color to her face. Her lips are a vibrant shade of pink, and her eyes are a darker shade of brown.",
@@ -127,6 +138,12 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             run_button = gr.Button("Generate as : ( 1280 x 832 )🤗", scale=0, elem_classes="submit-btn")
             
             with gr.Accordion("Advanced options", open=True, visible=True):
+                model_choice = gr.Dropdown(
+                    label="LoRA Selection",
+                    choices=list(LORA_OPTIONS.keys()),
+                    value="Sketch-Smudge"
+                )
+                
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -177,6 +194,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                     value=DEFAULT_STYLE_NAME,
                     label="Quality Style",
                 )
+
+
         
         with gr.Column(scale=2):
             result = gr.Gallery(label="Result", columns=1, show_label=False)
@@ -203,6 +222,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             randomize_seed,
             style_selection,
+            model_choice,
         ],
         outputs=[result, seed],
         api_name="run",