APG checkbox

app.py

@@ -12,6 +12,8 @@ from f_lite import FLitePipeline
 
 # Trick required because it is not a native diffusers model
 from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
+
+from f_lite.pipeline import APGConfig
 LOADABLE_CLASSES["f_lite"] = LOADABLE_CLASSES["f_lite.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
 ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]
 
@@ -26,7 +28,7 @@ else:
     logging.warning("GEMINI_API_KEY not found in environment variables. Prompt enrichment will not work.")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "Freepik/F-Lite"
+model_repo_id = "./grpo_hf"
 
 if torch.cuda.is_available():
     torch_dtype = torch.bfloat16
@@ -135,6 +137,7 @@ def infer(
     guidance_scale,
     num_inference_steps,
     use_prompt_enrichment,
+    enable_apg,
     progress=gr.Progress(track_tqdm=True),
 ):
     enriched_prompt_str = None
@@ -160,6 +163,7 @@ def infer(
         width=width,
         height=height,
         generator=generator,
+        apg_config=APGConfig(enabled=enable_apg)
     ).images[0]
     
     # Prepare Gradio updates for the enriched prompt display
@@ -287,6 +291,10 @@ with gr.Blocks(css=css, theme="ParityError/Interstellar") as demo:
                     step=0.1,
                     value=6,
                 )
+                enable_apg = gr.Checkbox(
+                    label="Enable APG",
+                    value=True,
+                )
 
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
@@ -334,6 +342,7 @@ with gr.Blocks(css=css, theme="ParityError/Interstellar") as demo:
             guidance_scale,
             num_inference_steps,
             use_prompt_enrichment,
+            enable_apg,
         ],
         outputs=[result, seed, enriched_prompt_display, enriched_prompt_text, enrichment_error],
     )
@@ -342,4 +351,4 @@ with gr.Blocks(css=css, theme="ParityError/Interstellar") as demo:
     gr.Markdown("[F-Lite Model Card and Weights](https://huggingface.co/Freepik/F-Lite)")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch() # server_name="0.0.0.0", share=True)