Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import gradio as gr
+from gradio_client import Client, handle_file
 import torch
 import spaces
 from diffusers import Lumina2Pipeline
@@ -10,17 +11,15 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
+
+def set_client_for_session(request: gr.Request):
+    x_ip_token = request.headers['x-ip-token']
+
+    # The "gradio/text-to-image" space is a ZeroGPU space
+    return Client("stzhao/LeX-Enhancer", headers={"X-IP-Token": x_ip_token})
+
 # Load models
-def load_models():
-    model_name = "X-ART/LeX-Enhancer-full"
-    
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="auto"
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
+def load_models():    
     pipe = Lumina2Pipeline.from_pretrained(
         "X-ART/LeX-Lumina", 
         torch_dtype=torch.bfloat16
@@ -28,52 +27,24 @@ def load_models():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     pipe.to("cuda")
     
-    return model, tokenizer, pipe
-
-model, tokenizer, pipe = load_models()
-
-def truncate_caption_by_tokens(caption, max_tokens=256):
-    """Truncate the caption to fit within the max token limit"""
-    tokens = tokenizer.encode(caption)
-    if len(tokens) > max_tokens:
-        truncated_tokens = tokens[:max_tokens]
-        caption = tokenizer.decode(truncated_tokens, skip_special_tokens=True)
-        print(f"Caption was truncated from {len(tokens)} tokens to {max_tokens} tokens")
-    return caption
-
-@spaces.GPU(duration=70)
-def generate_enhanced_caption(image_caption, text_caption):
-    # model.to("cuda")
-    """Generate enhanced caption using the LeX-Enhancer model"""
-    combined_caption = f"{image_caption}, with the text on it: {text_caption}."
-    instruction = """
-Below is the simple caption of an image with text. Please deduce the detailed description of the image based on this simple caption. Note: 1. The description should only include visual elements and should not contain any extended meanings. 2. The visual elements should be as rich as possible, such as the main objects in the image, their respective attributes, the spatial relationships between the objects, lighting and shadows, color style, any text in the image and its style, etc. 3. The output description should be a single paragraph and should not be structured. 4. The description should avoid certain situations, such as pure white or black backgrounds, blurry text, excessive rendering of text, or harsh visual styles. 5. The detailed caption should be human readable and fluent. 6. Avoid using vague expressions such as "may be" or "might be"; the generated caption must be in a definitive, narrative tone. 7. Do not use negative sentence structures, such as "there is nothing in the image," etc. The entire caption should directly describe the content of the image. 8. The entire output should be limited to 200 words.
-"""
-    messages = [
-        {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
-        {"role": "user", "content": instruction + "\nSimple Caption:\n" + combined_caption}
-    ]
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    return pipe
 
-    generated_ids = model.generate(
-        **model_inputs,
-        max_new_tokens=1024
-    )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
-
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    enhanced_caption = response.split("</think>", -1)[-1].strip(" ").strip("\n")
-    model.to("cpu")
-    torch.cuda.empty_cache()
-    
+def prompt_enhance(client, image_caption, text_caption):
+    combined_caption, enhanced_caption = client.predict(image_caption, text_caption, api_name="/generate_enhanced_caption")
     return combined_caption, enhanced_caption
+    
+
+pipe = load_models()
+
+# def truncate_caption_by_tokens(caption, max_tokens=256):
+#     """Truncate the caption to fit within the max token limit"""
+#     tokens = tokenizer.encode(caption)
+#     if len(tokens) > max_tokens:
+#         truncated_tokens = tokens[:max_tokens]
+#         caption = tokenizer.decode(truncated_tokens, skip_special_tokens=True)
+#         print(f"Caption was truncated from {len(tokens)} tokens to {max_tokens} tokens")
+#     return caption
+
 
 @spaces.GPU(duration=60)
 def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
@@ -81,7 +52,7 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
     pipe.enable_model_cpu_offload()
     """Generate image using LeX-Lumina"""
     # Truncate the caption if it's too long
-    enhanced_caption = truncate_caption_by_tokens(enhanced_caption, max_tokens=256)
+    # enhanced_caption = truncate_caption_by_tokens(enhanced_caption, max_tokens=256)
 
     print(f"enhanced caption:\n{enhanced_caption}")
     
@@ -107,13 +78,14 @@ def generate_image(enhanced_caption, seed, num_inference_steps, guidance_scale):
     
     return image
 
-@spaces.GPU(duration=130)
-def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer):
+# @spaces.GPU(duration=130)
+def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer, client):
     """Run the complete pipeline from captions to final image"""
     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)
+        # combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)
+        combined_caption, enhanced_caption = prompt_enhance(client, image_caption, text_caption)
     else:
         enhanced_caption = combined_caption
     
@@ -123,6 +95,7 @@ def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidanc
 
 # Gradio interface
 with gr.Blocks() as demo:
+    client = gr.State()
     gr.Markdown("# LeX-Enhancer & LeX-Lumina Demo")
     gr.Markdown("## Project Page: https://zhaoshitian.github.io/lexart/")
     gr.Markdown("Generate enhanced captions from simple image and text descriptions, then create images with LeX-Lumina")
@@ -208,9 +181,11 @@ with gr.Blocks() as demo:
     
     submit_btn.click(
         fn=run_pipeline,
-        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer],
+        inputs=[image_caption, text_caption, seed, num_inference_steps, guidance_scale, enable_enhancer, client],
         outputs=[output_image, combined_caption_box, enhanced_caption_box]
     )
 
+    demo.load(set_client_for_session, None, client)
+
 if __name__ == "__main__":
     demo.launch(debug=True)