Update app.py

app.py

@@ -1,3 +1,8 @@
+import gradio as gr
+import torch
+import spaces
+from src.models.transformer_sd3_SiamLayout import SiamLayoutSD3Transformer2DModel
+from src.pipeline.pipeline_CreatiLayout import CreatiLayoutSD3Pipeline
 from utils.bbox_visualization import bbox_visualization,scale_boxes
 from PIL import Image
 import os
@@ -5,6 +10,23 @@ import pandas as pd
 from huggingface_hub import login
 
 hf_token = os.getenv("HF_TOKEN")
+
+if hf_token is None:
+    raise ValueError("Hugging Face token not found. Please set the HF_TOKEN secret.")
+
+login(token=hf_token)
+
+model_path = "stabilityai/stable-diffusion-3-medium-diffusers"
+ckpt_path = "HuiZhang0812/CreatiLayout"
+
+transformer_additional_kwargs = dict(attention_type="layout",strict=True)
+
+transformer = SiamLayoutSD3Transformer2DModel.from_pretrained(
+     ckpt_path, subfolder="SiamLayout_SD3", torch_dtype=torch.float16,**transformer_additional_kwargs)
+
+pipe = CreatiLayoutSD3Pipeline.from_pretrained(model_path, transformer=transformer, torch_dtype=torch.float16)
+pipe = pipe.to("cuda")
+
 print("pipeline is loaded.")
 
 @spaces.GPU
@@ -12,6 +34,7 @@ def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame,
 
     if randomize_seed:
         seed = torch.randint(0, 100, (1,)).item()
+
     height = 1024
     width = 1024
 
@@ -19,7 +42,6 @@ def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame,
     box_detail_phrases_list_tmp = [c[0] for c in box_detail_phrases_list_tmp]
     boxes = boxes.astype(float).values.tolist()  
 
-
     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)
@@ -33,6 +55,7 @@ def process_image_and_text(global_caption, box_detail_phrases_list:pd.DataFrame,
         bbox_raw=boxes,
         height=height,
         width=width
+    ).images[0]
 
     return bbox_visualization_img, result_img
 
@@ -100,34 +123,12 @@ def get_samples():
 with gr.Blocks() 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)""")
-
-
-
-
-
-
-
     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.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)
@@ -137,7 +138,6 @@ with gr.Blocks() as demo:
             bbox_visualization_img = gr.Image(type="pil", label="Bounding Box Visualization")
         
         with gr.Column():
-
             output_image = gr.Image(type="pil", label="Generated Image")