Update app.py

app.py

@@ -1,347 +1,514 @@
-import os
-import uuid
-import gradio as gr
+import argparse
 import spaces
-from clip_slider_pipeline import CLIPSliderFlux
-from diffusers import FluxPipeline, AutoencoderTiny
+from visualcloze import VisualClozeModel
+import gradio as gr
+import examples
 import torch
-import numpy as np
-import cv2
-from PIL import Image
-from diffusers.utils import load_image
-from diffusers.utils import export_to_video
-import random
-
-# English menu labels
-english_labels = {
-    "Prompt": "Prompt",
-    "1st direction to steer": "1st Direction",
-    "2nd direction to steer": "2nd Direction",
-    "Strength": "Strength",
-    "Generate directions": "Generate Directions",
-    "Generated Images": "Generated Images",
-    "From 1st to 2nd direction": "From 1st to 2nd Direction",
-    "Strip": "Image Strip",
-    "Looping video": "Looping Video",
-    "Advanced options": "Advanced Options",
-    "Num of intermediate images": "Number of Intermediate Images",
-    "Num iterations for clip directions": "Number of CLIP Direction Iterations",
-    "Num inference steps": "Number of Inference Steps",
-    "Guidance scale": "Guidance Scale",
-    "Randomize seed": "Randomize Seed",
-    "Seed": "Seed"
-}
+from functools import partial
+from data.prefix_instruction import get_layout_instruction
+from huggingface_hub import snapshot_download
 
-# Load pipelines
-base_model = "black-forest-labs/FLUX.1-schnell"
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to("cuda")
-pipe = FluxPipeline.from_pretrained(
-    base_model,
-    vae=taef1,
-    torch_dtype=torch.bfloat16
-)
-pipe.transformer.to(memory_format=torch.channels_last)
-clip_slider = CLIPSliderFlux(pipe, device=torch.device("cuda"))
-MAX_SEED = 2**32 - 1
-
-def save_images_with_unique_filenames(image_list, save_directory):
-    if not os.path.exists(save_directory):
-        os.makedirs(save_directory)
-    paths = []
-    for image in image_list:
-        unique_filename = f"{uuid.uuid4()}.png"
-        file_path = os.path.join(save_directory, unique_filename)
-        image.save(file_path)
-        paths.append(file_path)
-    return paths
-
-def convert_to_centered_scale(num):
-    if num % 2 == 0:  # even
-        start = -(num // 2 - 1)
-        end = num // 2
-    else:  # odd
-        start = -(num // 2)
-        end = num // 2
-    return tuple(range(start, end + 1))
-
-def is_korean(text):
-    """한글 포함 여부 확인"""
-    return any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text)
-
-@spaces.GPU(duration=85)
-def generate(prompt,
-             concept_1,
-             concept_2,
-             scale,
-             randomize_seed=True,
-             seed=42,
-             recalc_directions=True,
-             iterations=200,
-             steps=3,
-             interm_steps=33,
-             guidance_scale=3.5,
-             x_concept_1="", x_concept_2="",
-             avg_diff_x=None,
-             total_images=[],
-             gradio_progress=gr.Progress()):
-    # Check if there is Korean text and warn if so
-    if is_korean(prompt) or is_korean(concept_1) or is_korean(concept_2):
-        print("Korean text detected. The model will use it directly without translation.")
-    
-    print(f"Prompt: {prompt}, ← {concept_2}, {concept_1} ➡️ . scale {scale}, interm steps {interm_steps}")
-    slider_x = [concept_2, concept_1]
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    if not sorted(slider_x) == sorted([x_concept_1, x_concept_2]) or recalc_directions:
-        gradio_progress(0, desc="Calculating directions...")
-        avg_diff = clip_slider.find_latent_direction(slider_x[0], slider_x[1], num_iterations=iterations)
-        x_concept_1, x_concept_2 = slider_x[0], slider_x[1]
-    else:
-        avg_diff = avg_diff_x
-    images = []
-    high_scale = scale
-    low_scale = -1 * scale
-    for i in gradio_progress.tqdm(range(interm_steps), desc="Generating images"):
-        cur_scale = low_scale + (high_scale - low_scale) * i / (interm_steps - 1)
-        image = clip_slider.generate(
-            prompt,
-            width=768,
-            height=768,
-            guidance_scale=guidance_scale,
-            scale=cur_scale,
-            seed=seed,
-            num_inference_steps=steps,
-            avg_diff=avg_diff
-        )
-        images.append(image)
-    canvas = Image.new('RGB', (256 * interm_steps, 256))
-    for i, im in enumerate(images):
-        canvas.paste(im.resize((256, 256)), (256 * i, 0))
-    comma_concepts_x = f"{slider_x[1]}, {slider_x[0]}"
-    scale_total = convert_to_centered_scale(interm_steps)
-    scale_min = scale_total[0]
-    scale_max = scale_total[-1]
-    scale_middle = scale_total.index(0)
-    post_generation_slider_update = gr.update(label=comma_concepts_x, value=0, minimum=scale_min, maximum=scale_max, interactive=True)
-    avg_diff_x = avg_diff.cpu()
-    video_path = f"{uuid.uuid4()}.mp4"
-    print(video_path)
-    return x_concept_1, x_concept_2, avg_diff_x, export_to_video(images, video_path, fps=5), canvas, images, images[scale_middle], post_generation_slider_update, seed
-
-def update_pre_generated_images(slider_value, total_images):
-    number_images = len(total_images) if total_images else 0
-    if number_images > 0:
-        scale_tuple = convert_to_centered_scale(number_images)
-        return total_images[scale_tuple.index(slider_value)][0]
-    else:
-        return None
-
-def reset_recalc_directions():
-    return True
-
-# Five "Time Stream" themed examples (one Korean example included)
-examples = [
-    ["신선한 토마토가 부패한 토마토로 변해가는 과정", "Fresh", "Rotten", 2.0],
-    ["A blooming flower gradually withers into decay", "Bloom", "Wither", 1.5],
-    ["A vibrant cityscape transforms into a derelict ruin over time", "Modern", "Ruined", 2.5],
-    ["A lively forest slowly changes into an autumnal landscape", "Spring", "Autumn", 2.0],
-    ["A calm ocean evolves into a stormy seascape as time passes", "Calm", "Stormy", 3.0]
-]
-
-css = """
-/* Bright and modern UI with background image */
-body {
-    background: #ffffff url('https://images.unsplash.com/photo-1506748686214-e9df14d4d9d0?ixlib=rb-1.2.1&auto=format&fit=crop&w=1600&q=80') no-repeat center center fixed;
-    background-size: cover;
-    font-family: "Helvetica Neue", Helvetica, Arial, sans-serif;
-    color: #333;
-}
-footer {
-    visibility: hidden;
-}
-.container {
-    max-width: 1200px;
-    margin: 20px auto;
-    padding: 0 10px;
-}
-.main-panel {
-    background-color: rgba(255, 255, 255, 0.9);
-    border-radius: 12px;
-    padding: 20px;
-    margin-bottom: 20px;
-    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
-}
-.controls-panel {
-    background-color: rgba(255, 255, 255, 0.85);
-    border-radius: 8px;
-    padding: 16px;
-    box-shadow: inset 0 2px 4px rgba(0, 0, 0, 0.05);
-}
-.image-display {
-    min-height: 400px;
-    display: flex;
-    flex-direction: column;
-    justify-content: center;
-}
-.slider-container {
-    padding: 10px 0;
-}
-.advanced-panel {
-    margin-top: 20px;
-    border-top: 1px solid #eaeaea;
-    padding-top: 20px;
+# Define the missing variables here
+GUIDANCE = """
+## How to use this demo:
+1. Select a task example from the right side, or prepare your own in-context examples and query.
+2. The grid will be filled with in-context examples and a query row.
+3. You can modify the task description or add content descriptions.
+4. Click "Generate" to create images following the pattern shown in examples.
+"""
+
+NOTE = """
+**Note:** The examples on the right side demonstrate various tasks. 
+Click on any example to load it into the interface. You can then modify images or prompts as needed.
+"""
+
+CITATION = """
+## Paper Citation
+```
+@article{liu2024visualcloze,
+  title={VisualCloze: A Universal Image Generation Framework via Visual In-Context Learning},
+  author={Liu, Zhaoyang and Lian, Yuheng and Wang, Jianfeng and Zhou, Aojun and Liu, Jiashi and Ye, Hang and Chen, Kai and Wang, Jingdong and Zhao, Deli},
+  journal={arXiv preprint arXiv:2504.07960},
+  year={2024}
 }
+```
 """
 
-# 여기서는 show_api=False를 Blocks(...)에 넣지 않음
-with gr.Blocks(css=css, title="Time Stream") as demo:
-    gr.Markdown("# Time Stream")
-    
-    x_concept_1 = gr.State("")
-    x_concept_2 = gr.State("")
-    total_images = gr.State([])
-    avg_diff_x = gr.State()
-    recalc_directions = gr.State(False)
-
-    with gr.Row(elem_classes="container"):
-        with gr.Column(scale=4):
-            with gr.Group(elem_classes="main-panel"):
-                gr.Markdown("### Image Generation Controls")
-                with gr.Group(elem_classes="controls-panel"):
-                    prompt = gr.Textbox(
-                        label=english_labels["Prompt"],
-                        info="Enter the description",
-                        placeholder="A dog in the park",
-                        lines=2
-                    )
-                    with gr.Row():
-                        with gr.Column(scale=1):
-                            concept_1 = gr.Textbox(
-                                label=english_labels["1st direction to steer"],
-                                info="Initial state",
-                                placeholder="Fresh"
-                            )
-                        with gr.Column(scale=1):
-                            concept_2 = gr.Textbox(
-                                label=english_labels["2nd direction to steer"],
-                                info="Final state",
-                                placeholder="Rotten"
+max_grid_h = 5
+max_grid_w = 5
+default_grid_h = 2
+default_grid_w = 3
+default_upsampling_noise = 0.4
+default_steps = 30
+
+
+def create_demo(model):
+    with gr.Blocks(title="VisualCloze Demo") as demo:
+        gr.Markdown("# VisualCloze: A Universal Image Generation Framework via Visual In-Context Learning")
+
+        gr.HTML("""
+        <div style="display:flex;column-gap:4px;">
+            <a href="https://github.com/lzyhha/VisualCloze">
+                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
+            </a> 
+            <a href="https://visualcloze.github.io/">
+                <img src='https://img.shields.io/badge/Project-Website-green'>
+            </a>
+            <a href="https://arxiv.org/abs/2504.07960">
+                <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
+            </a>
+            <a href="https://huggingface.co/VisualCloze/VisualCloze">
+                <img src='https://img.shields.io/badge/VisualCloze%20checkpoint-HF%20Model-green?logoColor=violet&label=%F0%9F%A4%97%20Checkpoint'>
+            </a>
+            <a href="https://huggingface.co/datasets/VisualCloze/Graph200K">
+                <img src='https://img.shields.io/badge/VisualCloze%20datasets-HF%20Dataset-6B88E3?logoColor=violet&label=%F0%9F%A4%97%20Graph200k%20Dataset'>
+            </a>
+        </div>
+        """)
+        
+        gr.Markdown(GUIDANCE)
+
+        # Pre-create all possible image components
+        all_image_inputs = []
+        rows = []
+        row_texts = []
+        with gr.Row():
+
+            with gr.Column(scale=2):
+                # Image grid
+                for i in range(max_grid_h):
+                    # Add row label before each row
+                    row_texts.append(gr.Markdown(
+                        "## Query" if i == default_grid_h - 1 else f"## In-context Example {i + 1}",  
+                        elem_id=f"row_text_{i}", 
+                        visible=i < default_grid_h
+                    ))
+                    with gr.Row(visible=i < default_grid_h, elem_id=f"row_{i}") as row:
+                        rows.append(row)
+                        for j in range(max_grid_w):
+                            img_input = gr.Image(
+                                label=f"In-context Example {i + 1}/{j + 1}" if i != default_grid_h - 1 else f"Query {j + 1}",
+                                type="pil",
+                                visible= i < default_grid_h and j < default_grid_w,
+                                interactive=True,
+                                elem_id=f"img_{i}_{j}"
                             )
-                    with gr.Row(elem_classes="slider-container"):
-                        x = gr.Slider(
-                            minimum=0,
-                            value=1.75,
-                            step=0.1,
-                            maximum=4.0,
-                            label=english_labels["Strength"],
-                            info="Maximum strength for each direction (above 2.5 may be unstable)"
-                        )
-                    submit = gr.Button(english_labels["Generate directions"], size="lg", variant="primary")
-            with gr.Accordion(label=english_labels["Advanced options"], open=False, elem_classes="advanced-panel"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        interm_steps = gr.Slider(
-                            label=english_labels["Num of intermediate images"],
-                            minimum=3,
-                            value=7,
-                            maximum=65,
-                            step=2
-                        )
-                    with gr.Column(scale=1):
-                        guidance_scale = gr.Slider(
-                            label=english_labels["Guidance scale"],
-                            minimum=0.1,
-                            maximum=10.0,
-                            step=0.1,
-                            value=3.5
-                        )
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        iterations = gr.Slider(
-                            label=english_labels["Num iterations for clip directions"],
-                            minimum=0,
-                            value=200,
-                            maximum=400,
-                            step=1
-                        )
-                    with gr.Column(scale=1):
-                        steps = gr.Slider(
-                            label=english_labels["Num inference steps"],
-                            minimum=1,
-                            value=3,
-                            maximum=4,
-                            step=1
-                        )
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        randomize_seed = gr.Checkbox(
-                            True,
-                            label=english_labels["Randomize seed"]
-                        )
-                    with gr.Column(scale=1):
-                        seed = gr.Slider(
-                            minimum=0,
-                            maximum=MAX_SEED,
-                            step=1,
-                            label=english_labels["Seed"],
-                            interactive=True,
-                            randomize=True
-                        )
-        # Right Column - Output
-        with gr.Column(scale=8):
-            with gr.Group(elem_classes="main-panel"):
-                gr.Markdown("### Generated Results")
-                image_strip = gr.Image(label="Image Strip", type="filepath", elem_id="strip", height=200)
-                output_video = gr.Video(label=english_labels["Looping video"], elem_id="video", loop=True, autoplay=True, height=600)
-                with gr.Row():
-                    post_generation_image = gr.Image(
-                        label=english_labels["Generated Images"],
-                        type="filepath",
-                        elem_id="interactive",
-                        elem_classes="image-display"
-                    )
-                    post_generation_slider = gr.Slider(
-                        minimum=-10,
-                        maximum=10,
-                        value=0,
-                        step=1,
-                        label=english_labels["From 1st to 2nd direction"]
+                            all_image_inputs.append(img_input)
+
+                # Prompts
+                layout_prompt = gr.Textbox(
+                    label="Layout Description (Auto-filled, Read-only)",
+                    placeholder="Layout description will be automatically filled based on grid size...",
+                    value=get_layout_instruction(default_grid_w, default_grid_h),
+                    elem_id="layout_prompt",
+                    interactive=False
+                )
+                
+                task_prompt = gr.Textbox(
+                    label="Task Description (Can be modified by referring to examples to perform custom tasks, but may lead to unstable results)",
+                    placeholder="Describe what task should be performed...",
+                    value="",
+                    elem_id="task_prompt"
+                )
+                
+                content_prompt = gr.Textbox(
+                    label="(Optional) Content Description (Image caption, Editing instructions, etc.)",
+                    placeholder="Describe the content requirements...",
+                    value="",
+                    elem_id="content_prompt"
+                )
+                
+                generate_btn = gr.Button("Generate", elem_id="generate_btn")
+                gr.Markdown(NOTE)
+
+                grid_h = gr.Slider(minimum=0, maximum=max_grid_h-1, value=default_grid_h-1, step=1, label="Number of In-context Examples", elem_id="grid_h")
+                grid_w = gr.Slider(minimum=1, maximum=max_grid_w, value=default_grid_w, step=1, label="Task Columns", elem_id="grid_w")
+                
+                with gr.Accordion("Advanced options", open=False):
+                    seed = gr.Number(label="Seed (0 for random)", value=0, precision=0)
+                    steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=default_steps, step=1)
+                    cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=50.0, value=30, step=1)
+                    upsampling_steps = gr.Slider(label="Upsampling steps (SDEdit)", minimum=1, maximum=100.0, value=10, step=1)
+                    upsampling_noise = gr.Slider(label="Upsampling noise (SDEdit)", minimum=0, maximum=1.0, value=default_upsampling_noise, step=0.05)
+
+                gr.Markdown(CITATION)
+
+            # Output
+            with gr.Column(scale=2):
+                output_gallery = gr.Gallery(
+                    label="Generated Results",
+                    show_label=True,
+                    elem_id="output_gallery",
+                    columns=None,
+                    rows=None,
+                    height="auto",
+                    allow_preview=True,
+                    object_fit="contain"
+                )
+            
+                gr.Markdown("# Task Examples")
+                gr.Markdown("Each click on a task may result in different examples.")
+                text_dense_prediction_tasks = gr.Textbox(label="Task", visible=False)
+                dense_prediction_tasks = gr.Dataset(
+                    samples=examples.dense_prediction_text, 
+                    label='Dense Prediction', 
+                    samples_per_page=1000, 
+                    components=[text_dense_prediction_tasks])
+                
+                text_conditional_generation_tasks = gr.Textbox(label="Task", visible=False)
+                conditional_generation_tasks = gr.Dataset(
+                    samples=examples.conditional_generation_text, 
+                    label='Conditional Generation',
+                    samples_per_page=1000, 
+                    components=[text_conditional_generation_tasks])
+                
+                text_image_restoration_tasks = gr.Textbox(label="Task", visible=False)
+                image_restoration_tasks = gr.Dataset(
+                    samples=examples.image_restoration_text, 
+                    label='Image Restoration',
+                    samples_per_page=1000, 
+                    components=[text_image_restoration_tasks])
+                
+                text_style_transfer_tasks = gr.Textbox(label="Task", visible=False)
+                style_transfer_tasks = gr.Dataset(
+                    samples=examples.style_transfer_text, 
+                    label='Style Transfer', 
+                    samples_per_page=1000, 
+                    components=[text_style_transfer_tasks])
+                
+                text_style_condition_fusion_tasks = gr.Textbox(label="Task", visible=False)
+                style_condition_fusion_tasks = gr.Dataset(
+                    samples=examples.style_condition_fusion_text, 
+                    label='Style Condition Fusion', 
+                    samples_per_page=1000, 
+                    components=[text_style_condition_fusion_tasks]) 
+                
+                text_tryon_tasks = gr.Textbox(label="Task", visible=False)
+                tryon_tasks = gr.Dataset(
+                    samples=examples.tryon_text, 
+                    label='Virtual Try-On', 
+                    samples_per_page=1000, 
+                    components=[text_tryon_tasks])
+
+                text_relighting_tasks = gr.Textbox(label="Task", visible=False)
+                relighting_tasks = gr.Dataset(
+                    samples=examples.relighting_text, 
+                    label='Relighting', 
+                    samples_per_page=1000, 
+                    components=[text_relighting_tasks])
+                
+                text_photodoodle_tasks = gr.Textbox(label="Task", visible=False)
+                photodoodle_tasks = gr.Dataset(
+                    samples=examples.photodoodle_text, 
+                    label='Photodoodle', 
+                    samples_per_page=1000, 
+                    components=[text_photodoodle_tasks])   
+
+                text_editing_tasks = gr.Textbox(label="Task", visible=False)
+                editing_tasks = gr.Dataset(
+                    samples=examples.editing_text, 
+                    label='Editing', 
+                    samples_per_page=1000, 
+                    components=[text_editing_tasks])
+
+                text_unseen_tasks = gr.Textbox(label="Task", visible=False)
+                unseen_tasks = gr.Dataset(
+                    samples=examples.unseen_tasks_text, 
+                    label='Unseen Tasks (May produce unstable effects)', 
+                    samples_per_page=1000, 
+                    components=[text_unseen_tasks]) 
+                
+                gr.Markdown("# Subject-driven Tasks Examples")
+                text_subject_driven_tasks = gr.Textbox(label="Task", visible=False)
+                subject_driven_tasks = gr.Dataset(
+                    samples=examples.subject_driven_text, 
+                    label='Subject-driven Generation', 
+                    samples_per_page=1000, 
+                    components=[text_subject_driven_tasks])
+                
+                text_condition_subject_fusion_tasks = gr.Textbox(label="Task", visible=False)
+                condition_subject_fusion_tasks = gr.Dataset(
+                    samples=examples.condition_subject_fusion_text, 
+                    label='Condition+Subject Fusion', 
+                    samples_per_page=1000, 
+                    components=[text_condition_subject_fusion_tasks])
+                
+                text_style_transfer_with_subject_tasks = gr.Textbox(label="Task", visible=False)
+                style_transfer_with_subject_tasks = gr.Dataset(
+                    samples=examples.style_transfer_with_subject_text, 
+                    label='Style Transfer with Subject', 
+                    samples_per_page=1000, 
+                    components=[text_style_transfer_with_subject_tasks])
+
+                text_condition_subject_style_fusion_tasks = gr.Textbox(label="Task", visible=False)
+                condition_subject_style_fusion_tasks = gr.Dataset(
+                    samples=examples.condition_subject_style_fusion_text, 
+                    label='Condition+Subject+Style Fusion', 
+                    samples_per_page=1000, 
+                    components=[text_condition_subject_style_fusion_tasks])
+
+                text_editing_with_subject_tasks = gr.Textbox(label="Task", visible=False)
+                editing_with_subject_tasks = gr.Dataset(
+                    samples=examples.editing_with_subject_text, 
+                    label='Editing with Subject', 
+                    samples_per_page=1000, 
+                    components=[text_editing_with_subject_tasks])
+
+                text_image_restoration_with_subject_tasks = gr.Textbox(label="Task", visible=False)
+                image_restoration_with_subject_tasks = gr.Dataset(
+                    samples=examples.image_restoration_with_subject_text, 
+                    label='Image Restoration with Subject', 
+                    samples_per_page=1000, 
+                    components=[text_image_restoration_with_subject_tasks])
+                
+        def update_grid(h, w):
+            actual_h = h + 1
+            model.set_grid_size(actual_h, w)
+            
+            updates = []
+            
+            # Update image component visibility
+            for i in range(max_grid_h * max_grid_w):
+                curr_row = i // max_grid_w
+                curr_col = i % max_grid_w
+                updates.append(
+                    gr.update(
+                        label=f"In-context Example {curr_row + 1}/{curr_col + 1}" if curr_row != actual_h - 1 else f"Query {curr_col + 1}",
+                        elem_id=f"img_{curr_row}_{curr_col}", 
+                        visible=(curr_row < actual_h and curr_col < w)))
+            
+            # Update row visibility and labels
+            updates_row = []
+            updates_row_text = []
+            for i in range(max_grid_h):
+                updates_row.append(gr.update(f"row_{i}", visible=(i < actual_h)))
+                updates_row_text.append(
+                    gr.update(
+                        elem_id=f"row_text_{i}", 
+                        visible=i < actual_h, 
+                        value="## Query" if i == actual_h - 1 else f"## In-context Example {i + 1}", 
                     )
+                )
+            
+            updates.extend(updates_row)
+            updates.extend(updates_row_text)
+            updates.append(gr.update(elem_id="layout_prompt", value=get_layout_instruction(w, actual_h)))
+            return updates
+
+        def generate_image(*inputs):
+            images = []
+            if grid_h.value + 1 != model.grid_h or grid_w.value != model.grid_w:
+                raise gr.Error('Please wait for the loading to complete.')
+            for i in range(model.grid_h):
+                images.append([])
+                for j in range(model.grid_w):
+                    images[i].append(inputs[i * max_grid_w + j])
+                    if i != model.grid_h - 1:
+                        if inputs[i * max_grid_w + j] is None:
+                            raise gr.Error('Please upload in-context examples. Possible that the task examples have not finished loading yet, and you can try waiting a few seconds before clicking the button again.')
+            seed, cfg, steps, upsampling_steps, upsampling_noise, layout_text, task_text, content_text = inputs[-8:]
+            
+            try:
+                results = generate(
+                    images, 
+                    [layout_text, task_text, content_text], 
+                    seed=seed, cfg=cfg, steps=steps, 
+                    upsampling_steps=upsampling_steps, upsampling_noise=upsampling_noise
+                )
+            except Exception as e:
+                raise gr.Error('Process error. Possible that the task examples have not finished loading yet, and you can try waiting a few seconds before clicking the button again. Error: ' + str(e))
+
+            output = gr.update(
+                elem_id='output_gallery', 
+                value=results, 
+                columns=min(len(results), 2),
+                rows=int(len(results) / 2 + 0.5))
+            
+            return output
+
+        def process_tasks(task, func):
+            outputs = func(task)
+            mask = outputs[0]
+            state = outputs[1:8]
+            if state[5] is None:
+                state[5] = default_upsampling_noise
+            if state[6] is None:
+                state[6] = default_steps
+            images = outputs[8:-len(mask)]
+            output = outputs[-len(mask):]
+            for i in range(len(mask)):
+                if mask[i] == 1:
+                    images.append(None)
+                else:
+                    images.append(output[-len(mask) + i])
+            
+            state[0] = state[0] - 1
+            cur_hrid_h = state[0]
+            cur_hrid_w = state[1]
+
+            current_example = [None] * 25
+            for i, image in enumerate(images):
+                pos = (i // cur_hrid_w) * 5 + (i % cur_hrid_w)
+                if image is not None:
+                    current_example[pos] = image
+            update_grid(cur_hrid_h, cur_hrid_w)
+            output = gr.update(
+                elem_id='output_gallery', 
+                value=[o for o, m in zip(output, mask) if m == 1], 
+                columns=min(sum(mask), 2),
+                rows=int(sum(mask) / 2 + 0.5))
+            return [output] + current_example + state
+        
+        dense_prediction_tasks.click(
+            partial(process_tasks, func=examples.process_dense_prediction_tasks), 
+            inputs=[dense_prediction_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full", 
+            show_progress_on=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps] + [generate_btn])
+
+        conditional_generation_tasks.click(
+            partial(process_tasks, func=examples.process_conditional_generation_tasks), 
+            inputs=[conditional_generation_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        image_restoration_tasks.click(
+            partial(process_tasks, func=examples.process_image_restoration_tasks), 
+            inputs=[image_restoration_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        style_transfer_tasks.click(
+            partial(process_tasks, func=examples.process_style_transfer_tasks), 
+            inputs=[style_transfer_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        style_condition_fusion_tasks.click(
+            partial(process_tasks, func=examples.process_style_condition_fusion_tasks), 
+            inputs=[style_condition_fusion_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+            
+        relighting_tasks.click(
+            partial(process_tasks, func=examples.process_relighting_tasks), 
+            inputs=[relighting_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        tryon_tasks.click(
+            partial(process_tasks, func=examples.process_tryon_tasks), 
+            inputs=[tryon_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+        
+        photodoodle_tasks.click(
+            partial(process_tasks, func=examples.process_photodoodle_tasks), 
+            inputs=[photodoodle_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        editing_tasks.click(
+            partial(process_tasks, func=examples.process_editing_tasks), 
+            inputs=[editing_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+        
+        unseen_tasks.click(
+            partial(process_tasks, func=examples.process_unseen_tasks), 
+            inputs=[unseen_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        subject_driven_tasks.click(
+            partial(process_tasks, func=examples.process_subject_driven_tasks), 
+            inputs=[subject_driven_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        style_transfer_with_subject_tasks.click(
+            partial(process_tasks, func=examples.process_style_transfer_with_subject_tasks), 
+            inputs=[style_transfer_with_subject_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        condition_subject_fusion_tasks.click(
+            partial(process_tasks, func=examples.process_condition_subject_fusion_tasks), 
+            inputs=[condition_subject_fusion_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        condition_subject_style_fusion_tasks.click(
+            partial(process_tasks, func=examples.process_condition_subject_style_fusion_tasks), 
+            inputs=[condition_subject_style_fusion_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        editing_with_subject_tasks.click(
+            partial(process_tasks, func=examples.process_editing_with_subject_tasks), 
+            inputs=[editing_with_subject_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+
+        image_restoration_with_subject_tasks.click(
+            partial(process_tasks, func=examples.process_image_restoration_with_subject_tasks), 
+            inputs=[image_restoration_with_subject_tasks], 
+            outputs=[output_gallery] + all_image_inputs + [grid_h, grid_w, layout_prompt, task_prompt, content_prompt, upsampling_noise, steps], 
+            show_progress="full")
+        # Initialize grid
+        model.set_grid_size(default_grid_h, default_grid_w)
+        
+        # Connect event processing function to all components that need updating
+        output_components = all_image_inputs + rows + row_texts + [layout_prompt]
+
+        grid_h.change(fn=update_grid, inputs=[grid_h, grid_w], outputs=output_components)
+        grid_w.change(fn=update_grid, inputs=[grid_h, grid_w], outputs=output_components)
+        
+        # Modify generate button click event
+        generate_btn.click(
+            fn=generate_image,
+            inputs=all_image_inputs + [seed, cfg, steps, upsampling_steps, upsampling_noise] + [layout_prompt, task_prompt, content_prompt],
+            outputs=output_gallery
+        )
+        
+    return demo
+
+
+@spaces.GPU()
+def generate(
+    images, 
+    prompts, 
+    seed, cfg, steps, 
+    upsampling_steps, upsampling_noise):
+    with torch.no_grad():
+        return model.process_images(
+            images=images, 
+            prompts=prompts, 
+            seed=seed, 
+            cfg=cfg, 
+            steps=steps, 
+            upsampling_steps=upsampling_steps, 
+            upsampling_noise=upsampling_noise)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default="checkpoints/visualcloze-384-lora.pth")
+    parser.add_argument("--precision", type=str, choices=["fp32", "bf16", "fp16"], default="bf16")
+    parser.add_argument("--resolution", type=int, default=384)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
     
-    gr.Examples(
-        examples=examples,
-        inputs=[prompt, concept_1, concept_2, x]
-    )
+    snapshot_download(repo_id="VisualCloze/VisualCloze", repo_type="model", local_dir="checkpoints")
+
+    # Initialize model
+    model = VisualClozeModel(resolution=args.resolution, model_path=args.model_path, precision=args.precision)
     
-    submit.click(
-        fn=generate,
-        inputs=[
-            prompt, concept_1, concept_2, x, randomize_seed, seed,
-            recalc_directions, iterations, steps, interm_steps,
-            guidance_scale, x_concept_1, x_concept_2, avg_diff_x, total_images
-        ],
-        outputs=[
-            x_concept_1, x_concept_2, avg_diff_x,
-            output_video,
-            image_strip,
-            total_images,
-            post_generation_image,
-            post_generation_slider,
-            seed
-        ]
-    )
+    # Create Gradio demo
+    demo = create_demo(model)
     
-    iterations.change(fn=reset_recalc_directions, outputs=[recalc_directions])
-    seed.change(fn=reset_recalc_directions, outputs=[recalc_directions])
-    post_generation_slider.change(
-        fn=update_pre_generated_images,
-        inputs=[post_generation_slider, total_images],
-        outputs=[post_generation_image],
-        queue=False,
-        show_progress="hidden",
-        concurrency_limit=None
-    )
-
-# demo.launch(...)에서만 show_api=False 설정
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    # Start Gradio server
+    demo.launch()