Update app.py

app.py

@@ -16,7 +16,7 @@ import torch
 from diffusers import FluxPipeline
 
 torch.backends.cuda.matmul.allow_tf32 = True
-
+loaded_acc = None
 class timer:
     def __init__(self, method_name="timed process"):
         self.method = method_name
@@ -31,7 +31,7 @@ if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
 
 pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.load_lora_weights(hf_hub_download("RED-AIGC/TDD", "FLUX.1-dev_tdd_lora_weights.safetensors"))
+pipe.load_lora_weights(hf_hub_download("RED-AIGC/TDD", "FLUX.1-dev_tdd_lora_weights.safetensors"),adapter_name="TDD")
 pipe.fuse_lora(lora_scale=0.125)
 pipe.to("cuda")
 
@@ -80,14 +80,37 @@ with gr.Blocks(css=css) as demo:
                         with gr.Row():
                             steps = gr.Slider(label="Inference Steps", minimum=4, maximum=10, step=1, value=8)
                             scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=3.5, step=0.1, value=2.0)
-                        
-                        seed = gr.Number(label="Seed", value=3420, precision=0)
+                        with gr.Row():
+                            seed = gr.Number(label="Seed", value=3420, precision=0)
+                        with gr.Row():
+                            acc = gr.Dropdown(
+                                label="Accelerate Lora",
+                                choices=["TDD", "TDD_adv"],
+                                value="TDD",
+                            )
                 
                 generate_btn = gr.Button("Generate Image", variant="primary", scale=1)
 
         with gr.Column(scale=4):
             output = gr.Image(label="Your Generated Image")
     
+    person="Portrait of a girl, high-end photography, prominent facial details, significant depth of field, soft lighting, late afternoon sun, in the style of Richard Avedon, Henri Cartier-Bresson, and Yousuf Karsh. Golden ratio composition, exquisite detail, Hasselblad H6D-100c, featured on the cover of photography magazines."
+    dog="Portrait photo of a Shiba Inu, photograph, highly detailed fur, warm and cheerful light, soft pastel tones, vibrant and sunny atmosphere, style by Tim Flach, bright daylight, natural setting, centered, extremely detailed, Nikon D850, award-winning photography"
+    # scenery="A high-detail shot of the sea and sky, capturing the vibrant blues and crisp sunlight. Shot with a Canon 5D Mark IV,  award-winning photography."
+    scenery="Capture the golden hour's warm embrace on a sandy beach with a Nikon D850, highlighting the vibrant blues and the soft, golden sunlight. This award-winning photography captures the essence of a summer sunset by the coast."
+    gr.Examples(
+        examples=[
+            [person, "TDD", 1024, 1024, 8, 2, 8888],
+            [dog, "TDD", 1024, 1024, 6, 2, 29],
+            [scenery, "TDD", 1024, 1024, 6, 2.2, 3777],
+        ],
+        # inputs=[prompt, negative_prompt, ckpt, acc, steps, guidance_scale, eta, seed],
+        inputs=[prompt,acc, height, width, steps, scales, seed],
+        outputs=output,
+        fn=process_image,
+        cache_examples="lazy",
+    )
+
     gr.Markdown(
         """
         <div style="max-width: 650px; margin: 2rem auto; padding: 1rem; border-radius: 10px; background-color: #f0f0f0;">
@@ -103,8 +126,14 @@ with gr.Blocks(css=css) as demo:
     )
 
     @spaces.GPU
-    def process_image(height, width, steps, scales, prompt, seed):
+    def process_image(prompt,acc,height, width, steps, scales, seed):
         global pipe
+        global loaded_acc
+        if loaded_acc != acc:
+            #pipe.load_lora_weights(ACC_lora[acc], adapter_name=acc)
+            pipe.set_adapters([acc], adapter_weights=[1.0])
+            print(pipe.get_active_adapters())
+            loaded_acc = acc
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
             return pipe(
                 prompt=[prompt],
@@ -118,7 +147,7 @@ with gr.Blocks(css=css) as demo:
 
     generate_btn.click(
         process_image,
-        inputs=[height, width, steps, scales, prompt, seed],
+        inputs=[prompt, acc,height, width, steps, scales, seed],
         outputs=output
     )