Update app.py

app.py

@@ -16,14 +16,14 @@ def load_models():
     
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch_bfloat16,
         # device_map="auto"
     )
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     
     pipe = Lumina2Pipeline.from_pretrained(
         "X-ART/LeX-Lumina", 
-        torch_dtype=torch.bfloat16
+        torch_dtype=torch_bfloat16
     )
     device = "cuda" if torch.cuda.is_available() else "cpu"
     # pipe.to("cuda")
@@ -107,9 +107,15 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
     return image
 
 @spaces.GPU(duration=100)
-def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale):
+def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer):
     """Run the complete pipeline from captions to final image"""
-    combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)
+    combined_caption = f"{image_caption}, with the text on it: {text_caption}."
+    
+    if enable_enhancer:
+        combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)
+    else:
+        enhanced_caption = combined_caption
+    
     image = generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale)
     
     return {
@@ -139,6 +145,11 @@ with gr.Blocks() as demo:
             )
             
             with gr.Accordion("Advanced Settings", open=False):
+                enable_enhancer = gr.Checkbox(
+                    label="Enable LeX-Enhancer",
+                    value=False,
+                    info="When enabled, the caption will be enhanced before image generation"
+                )
                 seed = gr.Slider(
                     minimum=0,
                     maximum=100000,
@@ -170,7 +181,7 @@ with gr.Blocks() as demo:
                 interactive=False
             )
             enhanced_caption_box = gr.Textbox(
-                label="Enhanced Caption",
+                label="Enhanced Caption" if enable_enhancer.value else "Final Caption",
                 interactive=False,
                 lines=5
             )
@@ -187,9 +198,19 @@ with gr.Blocks() as demo:
         label="Example Inputs"
     )
     
+    # Update the label of enhanced_caption_box based on checkbox state
+    def update_caption_label(enable_enhancer):
+        return gr.Textbox.update(label="Enhanced Caption" if enable_enhancer else "Final Caption")
+    
+    enable_enhancer.change(
+        fn=update_caption_label,
+        inputs=enable_enhancer,
+        outputs=enhanced_caption_box
+    )
+    
     submit_btn.click(
         fn=run_pipeline,
-        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale],
+        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer],
         outputs=[output_image, combined_caption_box, enhanced_caption_box]
     )