Update app.py

app.py

@@ -1,56 +1,44 @@
 import gradio as gr
 import numpy as np
 import random
-import os
 
 import spaces
 from pipeline_flux import FluxPipeline
 from transformer_flux import FluxTransformer2DModel
 import torch
 
-flux_model = "schnell"
+flux_model = "dev"
 bfl_repo = f"black-forest-labs/FLUX.1-{flux_model}"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
 
 transformer = FluxTransformer2DModel.from_pretrained(
-    bfl_repo, subfolder="transformer", torch_dtype=dtype
+    bfl_repo, subfolder="transformer", torch_dtype=dtype,
+    use_auth_token=os.getenv("HF_TOKEN"),
 )
 pipe = FluxPipeline.from_pretrained(bfl_repo, transformer=None, torch_dtype=dtype)
 pipe.transformer = transformer
 pipe.scheduler.config.use_dynamic_shifting = False
 pipe.scheduler.config.time_shift = 10
-pipe.enable_model_cpu_offload()
+# pipe.enable_model_cpu_offload()
 pipe = pipe.to(device)
 
-
-transformer2 = FluxTransformer2DModel.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", subfolder="transformer", torch_dtype=dtype,
-    use_auth_token=os.getenv("HF_TOKEN"),
-)
-pipe2 = FluxPipeline.from_pretrained(bfl_repo, transformer=None, torch_dtype=dtype)
-pipe2.transformer = transformer2
-pipe2.scheduler.config.use_dynamic_shifting = False
-pipe2.scheduler.config.time_shift = 10
-pipe2.enable_model_cpu_offload()
-pipe2 = pipe.to(device)
-
 pipe.load_lora_weights(
     "Huage001/URAE",
     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
@@ -64,17 +52,16 @@ def infer(
     width,
     height,
     num_inference_steps,
-    model='2k',
     progress=gr.Progress(track_tqdm=True),
+    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}")
-    pipe = pipe if model == "schnell" else pipe2
+    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 randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -138,14 +125,14 @@ with gr.Blocks(css=css) as demo:
 
             gr.Markdown("### Setting:")
 
-            model = gr.Radio(
-                label="Model",
-                choices=[
-                    ("FLUX.1 dev", "dev"),
-                    ("FLUX.1 schnell", "schnell"),
-                ],
-                value="2k",
-            )
+            # model = gr.Radio(
+            #     label="Model",
+            #     choices=[
+            #         ("2K model", "2k"),
+            #         ("4K model (beta)", "4k"),
+            #     ],
+            #     value="2k",
+            # )
 
             with gr.Row():
                 width = gr.Slider(
@@ -179,7 +166,7 @@ with gr.Blocks(css=css) as demo:
                 minimum=1,
                 maximum=50,
                 step=1,
-                value=4,  # Replace with defaults that work for your model
+                value=20,  # Replace with defaults that work for your model
             )
 
         with gr.Column(elem_id="col2"):
@@ -190,12 +177,12 @@ with gr.Blocks(css=css) as demo:
         fn=infer,
         inputs=[
             prompt,
+            # model,
             seed,
             randomize_seed,
             width,
             height,
             num_inference_steps,
-            model,
         ],
         outputs=[result, seed],
     )