Update app.py

app.py

@@ -14,12 +14,11 @@ from gradio_imageslider import ImageSlider
 
 translator = Translator()
 HF_TOKEN = os.environ.get("HF_TOKEN")
-basemodel = "black-forest-labs/FLUX.1-schnell"
 MAX_SEED = np.iinfo(np.int32).max
 CSS = "footer { visibility: hidden; }"
 JS = "function () { gradioURL = window.location.href; if (!gradioURL.endsWith('?__theme=dark')) { window.location.replace(gradioURL + '?__theme=dark'); } }"
 
-def enable_lora(lora_add): 
+def enable_lora(lora_add, basemodel): 
     return basemodel if not lora_add else lora_add
 
 async def generate_image(prompt, model, lora_word, width, height, scales, steps, seed):
@@ -34,8 +33,8 @@ async def generate_image(prompt, model, lora_word, width, height, scales, steps,
     except Exception as e: 
         raise gr.Error(f"Error en {e}")
 
-async def gen(prompt, lora_add, lora_word, width, height, scales, steps, seed, upscale_factor, process_upscale):
-    model = enable_lora(lora_add)
+async def gen(prompt, basemodel, lora_add, lora_word, width, height, scales, steps, seed, upscale_factor, process_upscale):
+    model = enable_lora(lora_add, basemodel)
     image, seed = await generate_image(prompt, model, lora_word, width, height, scales, steps, seed)
     image_path = "temp_image.png"
     image.save(image_path)
@@ -68,6 +67,7 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
             with gr.Column(scale=0.8):
                 with gr.Group():
                     prompt = gr.Textbox(label="Prompt")
+                    basemodel_choice = gr.Radio(label="Base Model", choices=["black-forest-labs/FLUX.1-schnell", "black-forest-labs/FLUX.1-DEV"], value="black-forest-labs/FLUX.1-schnell")
                     lora_add = gr.Textbox(label="Add Flux LoRA", info="Modelo Lora", lines=1, value="XLabs-AI/flux-RealismLora")
                     lora_word = gr.Textbox(label="Add Flux LoRA Trigger Word", info="Add the Trigger Word", lines=1, value="")
                     width = gr.Slider(label="Width", minimum=512, maximum=1280, step=8, value=512)
@@ -86,7 +86,7 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
             queue=False
         ).then(
             fn=gen,
-            inputs=[prompt, lora_add, lora_word, width, height, scales, steps, seed, upscale_factor, process_upscale],
+            inputs=[prompt, basemodel_choice, lora_add, lora_word, width, height, scales, steps, seed, upscale_factor, process_upscale],
             outputs=[output_res]
         )