Update app.py

app.py

@@ -25,7 +25,7 @@ pipe = pipe.to(device)
 
 
 transformer2 = FluxTransformer2DModel.from_pretrained(
-    bfl_repo, subfolder="transformer", torch_dtype=dtype
+    "black-forest-labs/FLUX.1-dev", subfolder="transformer", torch_dtype=dtype
 )
 pipe2 = FluxPipeline.from_pretrained(bfl_repo, transformer=None, torch_dtype=dtype)
 pipe2.transformer = transformer2
@@ -39,16 +39,16 @@ pipe.load_lora_weights(
     weight_name="urae_2k_adapter.safetensors",
     adapter_name="2k",
 )
-pipe.load_lora_weights(
-    "Huage001/URAE",
-    weight_name="urae_4k_adapter_lora_conversion_dev.safetensors",
-    adapter_name="4k_dev",
-)
-pipe.load_lora_weights(
-    "Huage001/URAE",
-    weight_name="urae_4k_adapter_lora_conversion_schnell.safetensors",
-    adapter_name="4k_schnell",
-)
+# pipe.load_lora_weights(
+#     "Huage001/URAE",
+#     weight_name="urae_4k_adapter_lora_conversion_dev.safetensors",
+#     adapter_name="4k_dev",
+# )
+# pipe.load_lora_weights(
+#     "Huage001/URAE",
+#     weight_name="urae_4k_adapter_lora_conversion_schnell.safetensors",
+#     adapter_name="4k_schnell",
+# )
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 USE_ZERO_GPU = True
@@ -66,12 +66,13 @@ def infer(
     model='2k',
 ):
     print("Using model:", model)
-    if model == "2k":
-        pipe.vae.enable_tiling(True)
-        pipe.set_adapters("2k")
-    elif model == "4k":
-        pipe.vae.enable_tiling(True)
-        pipe.set_adapters(f"4k_{flux_model}")
+    # if model == "2k":
+    # pipe.vae.enable_tiling(True)
+    # pipe.set_adapters("2k")
+    # # elif model == "4k":
+    #     pipe.vae.enable_tiling(True)
+    #     pipe.set_adapters(f"4k_{flux_model}")
+    pipe = pipe if model == "schnell" else pipe2
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -135,14 +136,14 @@ with gr.Blocks(css=css) as demo:
 
             gr.Markdown("### Setting:")
 
-            # model = gr.Radio(
-            #     label="Model",
-            #     choices=[
-            #         ("2K model", "2k"),
-            #         ("4K model (beta)", "4k"),
-            #     ],
-            #     value="2k",
-            # )
+            model = gr.Radio(
+                label="Model",
+                choices=[
+                    ("FLUX.1 dev", "dev"),
+                    ("FLUX.1 schnell", "schnell"),
+                ],
+                value="2k",
+            )
 
             with gr.Row():
                 width = gr.Slider(
@@ -187,7 +188,7 @@ with gr.Blocks(css=css) as demo:
         fn=infer,
         inputs=[
             prompt,
-            # model,
+            model,
             seed,
             randomize_seed,
             width,