Update app.py

app.py

@@ -35,8 +35,8 @@ def load_models():
 
 model, tokenizer, pipe = load_models()
 
-@spaces.GPU(duration=100)
-def generate_enhanced_caption(image_caption, text_caption):
+@spaces.GPU(duration=200)
+def generate_enhanced_caption(image_caption, text_caption, progress=gr.Progress(track_tqdm=True)):
     """Generate enhanced caption using the LeX-Enhancer model"""
     combined_caption = f"{image_caption}, with the text on it: {text_caption}."
     instruction = """
@@ -51,7 +51,8 @@ Below is the simple caption of an image with text. Please deduce the detailed de
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    # model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    model_inputs = tokenizer([text], return_tensors="pt")
 
     generated_ids = model.generate(
         **model_inputs,
@@ -68,8 +69,8 @@ Below is the simple caption of an image with text. Please deduce the detailed de
     
     return combined_caption, enhanced_caption
 
-@spaces.GPU(duration=100)
-def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
+@spaces.GPU(duration=200)
+def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale, progress=gr.Progress(track_tqdm=True)):
     """Generate image using LeX-Lumina"""
     generator = torch.Generator("cpu").manual_seed(seed) if seed != 0 else None
     
@@ -86,11 +87,11 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
         system_prompt="You are an assistant designed to generate superior images with the superior degree of image-text alignment based on textual prompts or user prompts.",
     ).images[0]
 
-    # torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
     
     return image
 
-# @spaces.GPU(duration=100)
+@spaces.GPU(duration=200)
 def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale):
     """Run the complete pipeline from captions to final image"""
     combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)