Update app.py

app.py

@@ -6,7 +6,6 @@ from src.pipeline.pipeline_CreatiLayout import CreatiLayoutSD3Pipeline
 from utils.bbox_visualization import bbox_visualization,scale_boxes
 from PIL import Image
 import os
-import pandas as pd
 from huggingface_hub import login
 
 hf_token = os.getenv("HF_TOKEN")
@@ -30,7 +29,7 @@ pipe = pipe.to("cuda")
 print("pipeline is loaded.")
 
 @spaces.GPU
-def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame, boxes:pd.DataFrame,seed: int=42, randomize_seed: bool=False, guidance_scale: float=7.5, num_inference_steps: int=50):
+def process_image_and_text(global_caption, region_captions_text, boxes_text, seed: int=42, randomize_seed: bool=False, guidance_scale: float=7.5, num_inference_steps: int=50):
 
     if randomize_seed:
         seed = torch.randint(0, 100, (1,)).item()
@@ -38,20 +37,37 @@ def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame,
     height = 1024
     width = 1024
 
-    box_detail_phrases_list_tmp = box_detail_phrases_list.values.tolist()
-    box_detail_phrases_list_tmp = [c[0] for c in box_detail_phrases_list_tmp]
-    boxes = boxes.astype(float).values.tolist()  
+    # 解析区域标题
+    region_captions = [line.strip() for line in region_captions_text.split('\n') if line.strip()]
+    
+    # 解析边界框
+    boxes = []
+    for line in boxes_text.split('\n'):
+        if line.strip():
+            try:
+                coords = line.split(',')
+                if len(coords) == 4:
+                    box = [float(coord.strip()) for coord in coords]
+                    boxes.append(box)
+            except ValueError:
+                continue
+
+    if not region_captions or not boxes:
+        raise gr.Error("Please provide both region captions and bounding boxes.")
+    
+    if len(region_captions) != len(boxes):
+        raise gr.Error("Number of region captions must match number of bounding boxes.")
 
-    white_image = Image.new('RGB', (width, height), color='rgb(256,256,256)')
-    show_input = {"boxes":scale_boxes(boxes,width,height),"labels":box_detail_phrases_list_tmp}
-    bbox_visualization_img = bbox_visualization(white_image,show_input)
+    white_image = Image.new('RGB', (width, height), color='white')
+    show_input = {"boxes": scale_boxes(boxes, width, height), "labels": region_captions}
+    bbox_visualization_img = bbox_visualization(white_image, show_input)
     
     result_img = pipe(
         prompt=global_caption,
         generator=torch.Generator(device="cuda").manual_seed(seed),
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-        bbox_phrases=box_detail_phrases_list_tmp,
+        bbox_phrases=region_captions,
         bbox_raw=boxes,
         height=height,
         width=width
@@ -59,109 +75,67 @@ def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame,
 
     return bbox_visualization_img, result_img
 
-def get_samples():
-    sample_list = [
-        {   
-            "global_caption": "A picturesque scene features Iron Man standing confidently on a rugged rock by the sea, holding a drawing board with his hands. The board displays the words 'Creative Layout' in a playful, hand-drawn font. The serene sea shimmers under the setting sun. The sky is painted with a gradient of warm colors, from deep oranges to soft purples.",
-            "region_caption_list": [
-                "Iron Man standing confidently on a rugged rock.",
-                "A rugged rock by the sea.",
-                "A drawing board with the words \"Creative Layout\" in a playful, hand-drawn font.",
-                "The serene sea shimmers under the setting sun.",
-                "The sky is a shade of deep orange to soft purple."
-            ],
-            "region_bboxes_list": [
-                [0.40, 0.35, 0.55, 0.80],  
-                [0.35, 0.75, 0.60, 0.95],  
-                [0.40, 0.45, 0.55, 0.65], 
-                [0.00, 0.30, 1.00, 0.90],  
-                [0.00, 0.00, 1.00, 0.30]
-            ]
-        },
-        {
-            "global_caption": "This is a photo showcasing two wooden benches in a park. The bench on the left is painted in a vibrant blue, while the one on the right is painted in a green. Both are placed on a path paved with stones, surrounded by lush trees and shrubs. The sunlight filters through the leaves, casting dappled shadows on the ground, creating a tranquil and comfortable atmosphere.",
-            "region_caption_list": [
-                "A weathered, blue wooden bench with green elements in a natural setting.",
-                "Old, weathered wooden benches with green paint.",
-                "A dirt path in a park with green grass on the sides and two colorful wooden benches.",
-                "Thick, verdant foliage of mature trees in a dense forest."
-            ],
-            "region_bboxes_list": [
-                [0.30, 0.44, 0.62, 0.78],  
-                [0.54, 0.41, 0.75, 0.65],  
-                [0.00, 0.39, 1.00, 1.00],  
-                [0.00, 0.00, 1.00, 0.43]  
-            ]
-        },
-        {
-            "global_caption": "This is a wedding photo taken in a photography studio, showing a newlywed couple sitting on a brown leather sofa in a modern indoor setting. The groom is dressed in a pink suit, paired with a pink tie and white shirt, while the bride is wearing a white wedding dress with a long veil. They are sitting on a brown leather sofa, with a wooden table in front of them, on which a bouquet of flowers is placed. The background is a bar with a staircase and a wall decorated with lights, creating a warm and romantic atmosphere.",
-            "region_caption_list": [
-                "A floral arrangement consisting of roses, carnations, and eucalyptus leaves on a wooden surface.",
-                "A white wedding dress with off-the-shoulder ruffles and a long, sheer veil.",
-                "A polished wooden table with visible grain and knots.",
-                "A close-up of a dark brown leather sofa with tufted upholstery and button details.",
-                "A man in a pink suit with a white shirt and red tie, sitting on a leather armchair.",
-                "A person in a suit seated on a leather armchair near a wooden staircase with books and bottles.",
-                "Bride in white gown with veil, groom in maroon suit and pink tie, seated on leather armchairs."
-            ],
-            "region_bboxes_list": [
-                [0.09, 0.65, 0.31, 0.93],
-                [0.62, 0.25, 0.89, 0.90],
-                [0.01, 0.70, 0.78, 0.99],
-                [0.76, 0.65, 1.00, 0.99],
-                [0.27, 0.32, 0.72, 0.75],
-                [0.00, 0.01, 0.52, 0.72],
-                [0.27, 0.09, 0.94, 0.89]
-            ]
-        }
-
-    ]
-    return [[sample["global_caption"], [[caption] for caption in sample["region_caption_list"]], sample["region_bboxes_list"]] for sample in sample_list]
-
-
-
-with gr.Blocks() as demo:
+# 示例数据
+example1_caption = "A picturesque scene features Iron Man standing confidently on a rugged rock by the sea, holding a drawing board with his hands. The board displays the words 'Creative Layout' in a playful, hand-drawn font. The serene sea shimmers under the setting sun. The sky is painted with a gradient of warm colors, from deep oranges to soft purples."
+example1_regions = """Iron Man standing confidently on a rugged rock.
+A rugged rock by the sea.
+A drawing board with the words "Creative Layout" in a playful, hand-drawn font.
+The serene sea shimmers under the setting sun.
+The sky is a shade of deep orange to soft purple."""
+example1_boxes = """0.40, 0.35, 0.55, 0.80
+0.35, 0.75, 0.60, 0.95
+0.40, 0.45, 0.55, 0.65
+0.00, 0.30, 1.00, 0.90
+0.00, 0.00, 1.00, 0.30"""
+
+with gr.Blocks(title="CreatiLayout") as demo:
     gr.Markdown("# CreatiLayout: Layout-to-Image generation")
-    gr.Markdown("""CreatiLayout is a layout-to-image framework for Diffusion Transformer models, offering high-quality and fine-grained controllable generation based on the global description and entity annotations. Users need to provide a global description and the position and description of each entity, as shown in the examples. Please feel free to modify the position and attributes of the entities in the examples (such as size, color, shape, text, portrait, etc.). Here are some inspirations: Iron Man -> Spider Man/Harry Potter/Buzz Lightyear; CreatiLayout -> Hello Friends/Let's Control; drawing board -> round drawing board; Modify the position of the drawing board to (0.4, 0.15, 0.55, 0.35)""")
+    gr.Markdown("""CreatiLayout is a layout-to-image framework for Diffusion Transformer models, offering high-quality and fine-grained controllable generation based on the global description and entity annotations. Users need to provide a global description and the position and description of each entity, as shown in the examples. Please feel free to modify the position and attributes of the entities in the examples (such as size, color, shape, text, portrait, etc.). 
+
+**Instructions:**
+- Enter a global caption describing the overall scene
+- List region captions (one per line) 
+- Enter bounding boxes as x1,y1,x2,y2 (one per line, values between 0 and 1)
+- Number of region captions must match number of bounding boxes""")
+    
     with gr.Row():
-        
-        with gr.Column():
-            global_caption = gr.Textbox(lines=2, label="Global Caption")
-            box_detail_phrases_list = gr.Dataframe(headers=["Region Captions"], label="Region Captions")
-            boxes = gr.Dataframe(headers=["x1", "y1", "x2", "y2"], label="Region Bounding Boxes (x_min,y_min,x_max,y_max)")
+        with gr.Column(scale=1):
+            global_caption = gr.Textbox(
+                lines=3, 
+                label="Global Caption",
+                placeholder="Describe the overall scene...",
+                value=example1_caption
+            )
+            region_captions_text = gr.Textbox(
+                lines=6, 
+                label="Region Captions (one per line)",
+                placeholder="Enter region descriptions, one per line...",
+                value=example1_regions
+            )
+            boxes_text = gr.Textbox(
+                lines=6, 
+                label="Region Bounding Boxes (x1,y1,x2,y2 per line)",
+                placeholder="0.1,0.1,0.5,0.5\n0.6,0.2,0.9,0.8",
+                value=example1_boxes
+            )
+            
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(0, 100, step=1, label="Seed", value=42)
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
                 guidance_scale = gr.Slider(1, 30, step=0.5, label="Guidance Scale", value=7.5)
                 num_inference_steps = gr.Slider(1, 50, step=1, label="Number of inference steps", value=28)
-        with gr.Column():
-            bbox_visualization_img = gr.Image(type="pil", label="Bounding Box Visualization")
+            
+            generate_btn = gr.Button("Generate", variant="primary", size="lg")
         
-        with gr.Column():
+        with gr.Column(scale=1):
+            bbox_visualization_img = gr.Image(type="pil", label="Bounding Box Visualization")
             output_image = gr.Image(type="pil", label="Generated Image")
 
-
-
-    gr.Button("Generate").click(
+    generate_btn.click(
         fn=process_image_and_text,
-        inputs=[global_caption, box_detail_phrases_list, boxes, seed, randomize_seed, guidance_scale, num_inference_steps],
+        inputs=[global_caption, region_captions_text, boxes_text, seed, randomize_seed, guidance_scale, num_inference_steps],
         outputs=[bbox_visualization_img, output_image]
     )
-    
-    
-    gr.Examples(
-        examples=get_samples(),
-        inputs=[global_caption, box_detail_phrases_list, boxes],
-        outputs=[bbox_visualization_img, output_image],
-        fn=process_image_and_text,
-        cache_examples=True
-    )
-    
-    
-    
-    
 
 if __name__ == "__main__":
-    demo.launch()
-
-
+    demo.launch()