Update app.py

app.py

@@ -1,154 +1,729 @@
 import gradio as gr
+import torch
+import base64
+import io
+from PIL import Image
+from transformers import (
+    LlavaNextProcessor, LlavaNextForConditionalGeneration,
+    T5EncoderModel, T5Tokenizer
+)
+from transformers import (
+    AutoProcessor, AutoModelForCausalLM, GenerationConfig,
+    T5EncoderModel, T5Tokenizer
+)
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler, FluxPipeline
+from tordi.diffusion.pipelines.onediffusion import OneDiffusionPipeline
+from tordi.models.denoiser.nextdit import NextDiT
+from tordi.dataset.utils import get_closest_ratio, ASPECT_RATIO_512
+from typing import List, Optional
+
+# Import additional libraries
+import matplotlib
 import numpy as np
-import random
+import cv2
 
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
+# Task-specific tokens
+TASK2SPECIAL_TOKENS = {
+    "text2image": "[[text2image]]",
+    "deblurring": "[[deblurring]]",
+    "inpainting": "[[image_inpainting]]",
+    "canny": "[[canny2image]]",
+    "super_resolution": "[[super_resolution]]",
+    "depth2image": "[[depth2image]]",
+    "hed2image": "[[hed2img]]",
+    "pose2image": "[[pose2image]]",
+    "semanticmap2image": "[[semanticmap2image]]",
+    "boundingbox2image": "[[boundingbox2image]]",
+    "image_editing": "[[image_editing]]",
+    "faceid": "[[faceid]]",
+    "multiview": "[[multiview]]",
+    "subject_driven": "[[subject_driven]]"
+}
+NEGATIVE_PROMPT = "monochrome, greyscale, low-res, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation"
+
+
+class LlavaCaptionProcessor:
+    def __init__(self):
+        model_name = "llava-hf/llama3-llava-next-8b-hf"
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        self.processor = LlavaNextProcessor.from_pretrained(model_name)
+        self.model = LlavaNextForConditionalGeneration.from_pretrained(
+            model_name, torch_dtype=dtype, low_cpu_mem_usage=True
+        ).to(device)
+        self.SPECIAL_TOKENS = "assistant\n\n\n"
+
+    def generate_response(self, image: Image.Image, msg: str) -> str:
+        conversation = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": msg}]}]
+        with torch.no_grad():
+            prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
+            inputs = self.processor(prompt, image, return_tensors="pt").to(self.model.device)
+            output = self.model.generate(**inputs, max_new_tokens=250)
+            response = self.processor.decode(output[0], skip_special_tokens=True)
+        return response.split(msg)[-1].strip()[len(self.SPECIAL_TOKENS):]
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
+    def process(self, images: List[Image.Image], msg: str = None) -> List[str]:
+        if msg is None:
+            msg = f"Describe the contents of the photo in 150 words or fewer."
+        try:
+            return [self.generate_response(img, msg) for img in images]
+        except Exception as e:
+            print(f"Error in process: {str(e)}")
+            raise
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+class MolmoCaptionProcessor:
+    def __init__(self):
+        pretrained_model_name = 'allenai/Molmo-7B-O-0924'
+        self.processor = AutoProcessor.from_pretrained(
+            pretrained_model_name,
+            trust_remote_code=True,
+            torch_dtype='auto',
+            device_map='auto'
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name,
+            trust_remote_code=True,
+            torch_dtype='auto',
+            device_map='auto'
+        )
 
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+    def generate_response(self, image: Image.Image, msg: str) -> str:
+        inputs = self.processor.process(
+            images=[image],
+            text=msg
+        )
+        # Move inputs to the correct device and make a batch of size 1
+        inputs = {k: v.to(self.model.device).unsqueeze(0) for k, v in inputs.items()}
+        
+        # Generate output
+        output = self.model.generate_from_batch(
+            inputs,
+            GenerationConfig(max_new_tokens=250, stop_strings="<|endoftext|>"),
+            tokenizer=self.processor.tokenizer
+        )
+        
+        # Only get generated tokens and decode them to text
+        generated_tokens = output[0, inputs['input_ids'].size(1):]
+        return self.processor.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
 
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
+    def process(self, images: List[Image.Image], msg: str = None) -> List[str]:
+        if msg is None:
+            msg = f"Describe the contents of the photo in 150 words or fewer."
+        try:
+            return [self.generate_response(img, msg) for img in images]
+        except Exception as e:
+            print(f"Error in process: {str(e)}")
+            raise
+
+def initialize_models():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = NextDiT.from_pretrained(
+        # "/data/input/duongl/finetuning_distributed_multiview_16x8_scalerays_dl3dv_dynamic_shift_softcap_editing/checkpoint-98000",
+        # "/data/input/duongl/data/input/duongl/finetuning_distributed_multiview_16x8_scalerays_dl3dv_dynamic_shift_softcap_trainingWithFluxScheduler/checkpoint-10000/",        # "lehduong/OneDiffusion",
+        "lehduong/OneDiffusion",
+        subfolder="transformer",
+        torch_dtype=torch.float32,
+    ).to(device)
+    vae = AutoencoderKL.from_pretrained("lehduong/OneDiffusion", subfolder="vae").to(device)
+    text_encoder = T5EncoderModel.from_pretrained("lehduong/OneDiffusion", subfolder="text_encoder", torch_dtype=torch.float16).to(device)
+    tokenizer = T5Tokenizer.from_pretrained("lehduong/OneDiffusion", subfolder="tokenizer")
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        # "stabilityai/stable-diffusion-3-medium-diffusers",
+        # "black-forest-labs/FLUX.1-dev",
+        "lehduong/OneDiffusion",
+        subfolder="scheduler"
+    )
+    # scheduler = FlowMatchEulerDiscreteScheduler(
+    #     base_image_seq_len=256,
+    #     base_shift=0.5,
+    #     max_image_seq_len=4096,
+    #     max_shift=1.16,
+    #     num_train_timesteps=1000,
+    #     shift=3.0,
+    #     use_dynamic_shifting=True
+    # )
+    pipeline = OneDiffusionPipeline(
+        vae=vae, text_encoder=text_encoder, transformer=model, tokenizer=tokenizer, scheduler=scheduler
+    ).to(torch.bfloat16)
+    molmo_caption_processor = MolmoCaptionProcessor() # LlavaCaptionProcessor()
+    return pipeline, molmo_caption_processor
+
+def colorize_depth_maps(
+    depth_map, min_depth, max_depth, cmap="Spectral", valid_mask=None
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
+    """
+    Colorize depth maps with reversed colors.
+    """
+    assert len(depth_map.shape) >= 2, "Invalid dimension"
+    
+    if isinstance(depth_map, torch.Tensor):
+        depth = depth_map.detach().squeeze().numpy()
+    elif isinstance(depth_map, np.ndarray):
+        depth = depth_map.copy().squeeze()
+    # reshape to [ (B,) H, W ]
+    if depth.ndim < 3:
+        depth = depth[np.newaxis, :, :]
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+    # Normalize depth values to [0, 1]
+    depth = ((depth - min_depth) / (max_depth - min_depth)).clip(0, 1)
+    # Invert the depth values to reverse the colors
+    depth = 1 - depth
 
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+    # Use the colormap
+    cm = matplotlib.colormaps[cmap]
+    img_colored_np = cm(depth, bytes=False)[:, :, :, 0:3]  # values from 0 to 1
+    img_colored_np = np.rollaxis(img_colored_np, 3, 1)
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+    if valid_mask is not None:
+        if isinstance(depth_map, torch.Tensor):
+            valid_mask = valid_mask.detach().numpy()
+        valid_mask = valid_mask.squeeze()  # [H, W] or [B, H, W]
+        if valid_mask.ndim < 3:
+            valid_mask = valid_mask[np.newaxis, np.newaxis, :, :]
+        else:
+            valid_mask = valid_mask[:, np.newaxis, :, :]
+        valid_mask = np.repeat(valid_mask, 3, axis=1)
+        img_colored_np[~valid_mask] = 0
 
-        result = gr.Image(label="Result", show_label=False)
+    if isinstance(depth_map, torch.Tensor):
+        img_colored = torch.from_numpy(img_colored_np).float()
+    elif isinstance(depth_map, np.ndarray):
+        img_colored = img_colored_np
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
+    return img_colored
+
+
+def format_prompt(task_type: str, captions: List[str]) -> str:
+    if not captions:
+        return ""
+    if task_type == "faceid":
+        img_prompts = [f"[[img{i}]] {caption}" for i, caption in enumerate(captions, start=1)]
+        return f"[[faceid]] [[img0]] insert/your/caption/here {' '.join(img_prompts)}"
+    elif task_type == "image_editing":
+        target_caption = captions[0] if len(captions) >= 2 else "Insert target caption here"
+        source_caption = captions[1] if len(captions) >= 2 else captions[0]
+        return f"[[image_editing]] [[target_caption]] {target_caption} [[source_caption]] {source_caption}"
+    elif task_type == "semanticmap2image":
+        return f"[[semanticmap2image]] <#00ffff Cyan mask: insert/concept/to/segment/here> {captions[0]}"
+    elif task_type == "boundingbox2image":
+        return f"[[boundingbox2image]] <#00ffff Cyan boundingbox: insert/concept/to/segment/here> {captions[0]}"
+    elif task_type == "multiview":
+        # img_prompts = [f"[[img{i}]] {caption}" for i, caption in enumerate(captions)]
+        img_prompts = captions[0]
+        return f"[[multiview]] {img_prompts}"
+    elif task_type == "subject_driven":
+        return f"[[subject_driven]] <item: insert/item/here> [[img0]] insert/your/target/caption/here [[img1]] {captions[0]}"
+    else:
+        return f"{TASK2SPECIAL_TOKENS[task_type]} {captions[0]}"
+
+def update_prompt(images: List[Image.Image], task_type: str, custom_msg: str = None):
+    if not images:
+        return format_prompt(task_type, []), "Please upload at least one image!"
+    try:
+        captions = molmo_processor.process(images, custom_msg)
+        if not captions:
+            return "", "No valid images found!"
+        prompt = format_prompt(task_type, captions)
+        return prompt, f"Generated {len(captions)} captions successfully!"
+    except Exception as e:
+        return "", f"Error generating captions: {str(e)}"
+
+def generate_image(images: List[Image.Image], prompt: str, negative_prompt: str, num_inference_steps: int, guidance_scale: float, pag_guidance_scale: float, 
+                   denoise_mask: List[str], task_type: str, azimuth: str, elevation: str, distance: str, focal_length: float,
+                   height: int = 1024, width: int = 1024, scale_factor: float = 1.0, scale_watershed: float = 1.0,
+                   noise_scale: float = None, progress=gr.Progress()):
+    try:
+        img2img_kwargs = {
+            'prompt': prompt,
+            'negative_prompt': negative_prompt,
+            'num_inference_steps': num_inference_steps,
+            'guidance_scale': guidance_scale,
+            'height': height,
+            'width': width,
+            'forward_kwargs': {
+                'scale_factor': scale_factor,
+                'scale_watershed': scale_watershed
+            },
+            'noise_scale': noise_scale  # Added noise_scale here
+        }
+
+        if task_type == 'multiview':
+            # Parse azimuth, elevation, and distance into lists, allowing 'None' values
+            azimuths = [float(a.strip()) if a.strip().lower() != 'none' else None for a in azimuth.split(',')] if azimuth else []
+            elevations = [float(e.strip()) if e.strip().lower() != 'none' else None for e in elevation.split(',')] if elevation else []
+            distances = [float(d.strip()) if d.strip().lower() != 'none' else None for d in distance.split(',')] if distance else []
+
+            num_views = max(len(images), len(azimuths), len(elevations), len(distances))
+            if num_views == 0:
+                return None, "At least one image or camera parameter must be provided."
+
+            total_components = []
+            for i in range(num_views):
+                total_components.append(f"image_{i}")
+                total_components.append(f"camera_pose_{i}")
+
+            denoise_mask_int = [1 if comp in denoise_mask else 0 for comp in total_components]
+
+            if len(denoise_mask_int) != len(total_components):
+                return None, f"Denoise mask length mismatch: expected {len(total_components)} components."
+
+            # Pad the input lists to num_views length
+            images_padded = images + [] * (num_views - len(images))  # Do not add None
+            azimuths_padded = azimuths + [None] * (num_views - len(azimuths))
+            elevations_padded = elevations + [None] * (num_views - len(elevations))
+            distances_padded = distances + [None] * (num_views - len(distances))
+
+            # Prepare values
+            img2img_kwargs.update({
+                'image': images_padded,
+                'multiview_azimuths': azimuths_padded,
+                'multiview_elevations': elevations_padded,
+                'multiview_distances': distances_padded,
+                'multiview_focal_length': focal_length,  # Pass focal_length here
+                'is_multiview': True,
+                'denoise_mask': denoise_mask_int,
+                # 'predict_camera_poses': True,
+            })
+        else:
+            total_components = ["image_0"] + [f"image_{i+1}" for i in range(len(images))]
+            denoise_mask_int = [1 if comp in denoise_mask else 0 for comp in total_components]
+            if len(denoise_mask_int) != len(total_components):
+                return None, f"Denoise mask length mismatch: expected {len(total_components)} components."
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+            img2img_kwargs.update({
+                'image': images,
+                'denoise_mask': denoise_mask_int
+            })
+
+        progress(0, desc="Generating image...")
+        if task_type == 'text2image':
+            output = pipeline(
+                prompt=prompt, 
+                negative_prompt=negative_prompt, 
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                pag_guidance_scale=pag_guidance_scale,
+                height=height, 
+                width=width,
+                scale_factor=scale_factor,
+                scale_watershed=scale_watershed,
+                noise_scale=noise_scale  # Added noise_scale here
             )
+        else:
+            output = pipeline.img2img(**img2img_kwargs)
+        progress(1, desc="Done!")
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        # Process the output images if task is 'depth2image' and predicting depth
+        if task_type == 'depth2image' and denoise_mask_int[-1] == 1:
+            processed_images = []
+            for img in output.images:
+                depth_map = np.array(img.convert('L'))  # Convert to grayscale numpy array
+                min_depth = depth_map.min()
+                max_depth = depth_map.max()
+                colorized = colorize_depth_maps(depth_map, min_depth, max_depth)[0]
+                colorized = np.transpose(colorized, (1, 2, 0))
+                colorized = (colorized * 255).astype(np.uint8)
+                img_colorized = Image.fromarray(colorized)
+                processed_images.append(img_colorized)
+            output_images = processed_images + output.images
+        elif task_type in ['boundingbox2image', 'semanticmap2image'] and denoise_mask_int == [0,1] and images:
+            # Interpolate between input and output images
+            processed_images = []
+            for input_img, output_img in zip(images, output.images):
+                input_img_resized = input_img.resize(output_img.size)
+                blended_img = Image.blend(input_img_resized, output_img, alpha=0.5)
+                processed_images.append(blended_img)
+            output_images = processed_images + output.images
+        else:
+            output_images = output.images
 
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
+        return output_images, "Generation completed successfully!"
 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
+    except Exception as e:
+        return None, f"Error during generation: {str(e)}"
+
+def update_denoise_checkboxes(images_state: List[Image.Image], task_type: str, azimuth: str, elevation: str, distance: str):
+    if task_type == 'multiview':
+        azimuths = [a.strip() for a in azimuth.split(',')] if azimuth else []
+        elevations = [e.strip() for e in elevation.split(',')] if elevation else []
+        distances = [d.strip() for d in distance.split(',')] if distance else []
+        images_len = len(images_state)
+
+        num_views = max(images_len, len(azimuths), len(elevations), len(distances))
+        if num_views == 0:
+            return gr.update(choices=[], value=[]), "Please provide at least one image or camera parameter."
+
+        # Pad lists to the same length
+        azimuths += ['None'] * (num_views - len(azimuths))
+        elevations += ['None'] * (num_views - len(elevations))
+        distances += ['None'] * (num_views - len(distances))
+        # Do not add None to images_state
+
+        labels = []
+        values = []
+        for i in range(num_views):
+            labels.append(f"image_{i}")
+            labels.append(f"camera_pose_{i}")
+
+            # Default behavior: condition on provided inputs, generate missing ones
+            if i >= images_len:
+                values.append(f"image_{i}")
+            if azimuths[i].lower() == 'none' or elevations[i].lower() == 'none' or distances[i].lower() == 'none':
+                values.append(f"camera_pose_{i}")
+
+        return gr.update(choices=labels, value=values)
+    else:
+        labels = ["image_0"] + [f"image_{i+1}" for i in range(len(images_state))]
+        values = ["image_0"]
+        return gr.update(choices=labels, value=values)
+
+def apply_mask(images_state):
+    if len(images_state) < 2:
+        return None, "Please upload at least two images: first as the base image, second as the mask."
+    base_img = images_state[0]
+    mask_img = images_state[1]
+
+    # Convert images to arrays
+    base_arr = np.array(base_img)
+    mask_arr = np.array(mask_img)
+
+    # Convert mask to grayscale
+    if mask_arr.ndim == 3:
+        gray_mask = cv2.cvtColor(mask_arr, cv2.COLOR_RGB2GRAY)
+    else:
+        gray_mask = mask_arr
+
+    # Create a binary mask where non-black pixels are True
+    binary_mask = gray_mask > 10
+
+    # Define the gray color
+    gray_color = np.array([128, 128, 128], dtype=np.uint8)
+
+    # Apply gray color where mask is True
+    masked_arr = base_arr.copy()
+    masked_arr[binary_mask] = gray_color
+
+    masked_img = Image.fromarray(masked_arr)
+    return [masked_img], "Mask applied successfully!"
+
+def process_images_for_task_type(images_state: List[Image.Image], task_type: str):
+    # No changes needed here since we are processing the output images
+    return images_state, images_state
+
+# Initialize models
+pipeline, molmo_processor = initialize_models()
+
+with gr.Blocks(title="OneDiffusion Demo") as demo:
+    gr.Markdown("""
+    # OneDiffusion Demo
+
+    **Welcome to the OneDiffusion Demo!**
+
+    This application allows you to generate images based on your input prompts for various tasks. Here's how to use it:
+
+    1. **Select Task Type**: Choose the type of task you want to perform from the "Task Type" dropdown menu.
+
+    2. **Upload Images**: Drag and drop images directly onto the upload area, or click to select files from your device.
 
+    3. **Generate Captions**: **If you upload any images**, Click the "Generate Captions with Molmo" button to generate descriptive captions for your uploaded images (depend on the task). You can enter a custom message in the "Custom Message for Molmo" textbox e.g., "caption in 50 words" instead of 100 words.
+
+    4. **Configure Generation Settings**: Expand the "Advanced Configuration" section to adjust parameters like the number of inference steps, guidance scale, image size, and more.
+
+    5. **Generate Images**: After setting your preferences, click the "Generate Image" button. The generated images will appear in the "Generated Images" gallery.
+
+    6. **Manage Images**: Use the "Delete Selected Images" or "Delete All Images" buttons to remove unwanted images from the gallery.
+
+    **Notes**:
+    - For text-to-image:
+        + simply enter your prompt in this format "[[text2image]] your/prompt/here" and press the "Generate Image" button.
+        
+    - For boundingbox2image/semantic2image/inpainting etc tasks:
+        + To perform condition-to-image such as semantic map to image, follow above steps
+        + For image-to-condition e.g., image to depth, change the denoise_mask checkbox before generating images. You must UNCHECK image_0 box and CHECK image_1 box.
+        
+    - For FaceID tasks: 
+        + Use 3 or 4 images if single input image does not give satisfactory results.
+        + All images will be resized and center cropped to the input height and width. You should choose height and width so that faces in input images won't be cropped.
+        + Model works best with close-up portrait (input and output) images.
+        + If the model does not conform your text prompt, try using shorter caption for source image(s).
+        + If you have non-human subjects and does not get satisfactory results, try "copying" part of caption of source images where it describes the properties of the subject e.g., a monster with red eyes, sharp teeth, etc.
+        
+    - For Multiview generation:
+        + Only support square images (ideally in 512x512 resolution).
+        + Ensure the number of elevations, azimuths, and distances are equal. 
+        + The model generally works well for 2-5 views (include both input and generated images). Since the model is trained with 3 views on 512x512 resolution, you might try scale_factor of [1.1; 1.5] and scale_watershed of [100; 400] for better extrapolation.
+        + For better results:
+            1) try increasing num_inference_steps to 75-100.
+            2) avoid aggressively changes in target camera poses, for example to generate novel views at azimuth of 180, (simultaneously) generate 4 views with azimuth of 45, 90, 135, 180.
+    
+    Enjoy creating images with OneDiffusion!
+    """)
+
+    with gr.Row():
+        with gr.Column():
+            images_state = gr.State([])
+            selected_indices_state = gr.State([])
+            
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                gallery = gr.Gallery(
+                    label="Input Images",
+                    show_label=True,
+                    columns=2,
+                    rows=2,
+                    height="auto",
+                    object_fit="contain"
                 )
+            
+            # In the UI section, update the file_output component:
+            file_output = gr.File(
+                file_count="multiple",
+                file_types=["image"],
+                label="Drag and drop images here or click to upload",
+                height=100,
+                scale=2,
+                type="filepath"  # Add this parameter
+            )
+            
+            with gr.Row():
+                delete_button = gr.Button("Delete Selected Images")
+                delete_all_button = gr.Button("Delete All Images")
+            
+            task_type = gr.Dropdown(
+                choices=list(TASK2SPECIAL_TOKENS.keys()),
+                value="text2image",
+                label="Task Type"
+            )
+            
+            molmo_message = gr.Textbox(
+                lines=2,
+                value="Describe the contents of the photo in 100 words.",
+                label="Custom message for Molmo captioner"
+            )
+            
+            auto_caption_btn = gr.Button("Generate Captions with Molmo")
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
+        with gr.Column():
+            prompt = gr.Textbox(
+                lines=3,
+                placeholder="Enter your prompt here or use auto-caption...",
+                label="Prompt"
+            )
+            negative_prompt = gr.Textbox(
+                lines=3,
+                value=NEGATIVE_PROMPT,
+                placeholder="Enter negative prompt here...",
+                label="Negative Prompt"
+            )
+            caption_status = gr.Textbox(label="Caption Status")
+            
+    num_steps = gr.Slider(
+        minimum=1,
+        maximum=200,
+        value=30,
+        step=1,
+        label="Number of Inference Steps"
+    )
+    guidance_scale = gr.Slider(
+        minimum=0.1,
+        maximum=10.0,
+        value=4,
+        step=0.1,
+        label="Guidance Scale"
+    )
+    pag_guidance_scale = gr.Slider(
+        minimum=0.1,
+        maximum=10.0,
+        value=1,
+        step=0.1,
+        label="PAG guidance Scale"
+    )
+    height = gr.Number(value=1024, label="Height")
+    width = gr.Number(value=1024, label="Width")
+    
+    with gr.Accordion("Advanced Configuration", open=False):
+        with gr.Row():
+            denoise_mask_checkbox = gr.CheckboxGroup(
+                label="Denoise Mask",
+                choices=["image_0"],
+                value=["image_0"]
+            )
+            azimuth = gr.Textbox(
+                value="0",
+                label="Azimuths (degrees, comma-separated, 'None' for missing)"
+            )
+            elevation = gr.Textbox(
+                value="0",
+                label="Elevations (degrees, comma-separated, 'None' for missing)"
+            )
+            distance = gr.Textbox(
+                value="1.5",
+                label="Distances (comma-separated, 'None' for missing)"
+            )
+            focal_length = gr.Number(
+                value=1.3887,
+                label="Focal Length of camera for multiview generation"
+            )
+            scale_factor = gr.Number(value=1.0, label="Scale Factor")
+            scale_watershed = gr.Number(value=1.0, label="Scale Watershed")
+            noise_scale = gr.Number(value=1.0, label="Noise Scale")  # Added noise_scale input
+
+    output_images = gr.Gallery(
+        label="Generated Images",
+        show_label=True,
+        columns=4,
+        rows=2,
+        height="auto",
+        object_fit="contain"
+    )
+    
+    with gr.Column():
+        generate_btn = gr.Button("Generate Image")
+        apply_mask_btn = gr.Button("Apply Mask")
+    
+    status = gr.Textbox(label="Generation Status")
+
+    # Event Handlers
+    def update_gallery(files, images_state):
+        if not files:
+            return images_state, images_state
+        
+        new_images = []
+        for file in files:
+            try:
+                # Handle both file paths and file objects
+                if isinstance(file, dict):  # For drag and drop files
+                    file = file['path']
+                elif hasattr(file, 'name'):  # For uploaded files
+                    file = file.name
+                    
+                img = Image.open(file).convert('RGB')
+                new_images.append(img)
+            except Exception as e:
+                print(f"Error loading image: {str(e)}")
+                continue
+                
+        images_state.extend(new_images)
+        return images_state, images_state
 
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+    def on_image_select(evt: gr.SelectData, selected_indices_state):
+        selected_indices = selected_indices_state or []
+        index = evt.index
+        if index in selected_indices:
+            selected_indices.remove(index)
+        else:
+            selected_indices.append(index)
+        return selected_indices
+
+    def delete_images(selected_indices, images_state):
+        updated_images = [img for i, img in enumerate(images_state) if i not in selected_indices]
+        selected_indices_state = []
+        return updated_images, updated_images, selected_indices_state
+
+    def delete_all_images(images_state):
+        updated_images = []
+        selected_indices_state = []
+        return updated_images, updated_images, selected_indices_state
+
+    def update_height_width(images_state):
+        if images_state:
+            closest_ar = get_closest_ratio(
+                height=images_state[0].size[1],
+                width=images_state[0].size[0],
+                ratios=ASPECT_RATIO_512
+            )
+            height_val, width_val = int(closest_ar[0][0]), int(closest_ar[0][1])
+        else:
+            height_val, width_val = 1024, 1024  # Default values
+        return gr.update(value=height_val), gr.update(value=width_val)
+
+    # Connect events
+    file_output.change(
+        fn=update_gallery,
+        inputs=[file_output, images_state],
+        outputs=[images_state, gallery]
+    ).then(
+        fn=update_height_width,
+        inputs=[images_state],
+        outputs=[height, width]
+    ).then(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    gallery.select(
+        fn=on_image_select,
+        inputs=[selected_indices_state],
+        outputs=[selected_indices_state]
+    )
+
+    delete_button.click(
+        fn=delete_images,
+        inputs=[selected_indices_state, images_state],
+        outputs=[images_state, gallery, selected_indices_state]
+    ).then(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    delete_all_button.click(
+        fn=delete_all_images,
+        inputs=[images_state],
+        outputs=[images_state, gallery, selected_indices_state]
+    ).then(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    task_type.change(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    azimuth.change(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    elevation.change(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    distance.change(
+        fn=update_denoise_checkboxes,
+        inputs=[images_state, task_type, azimuth, elevation, distance],
+        outputs=[denoise_mask_checkbox]
+    )
+
+    generate_btn.click(
+        fn=generate_image,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
+            images_state, prompt, negative_prompt, num_steps, guidance_scale, pag_guidance_scale,
+            denoise_mask_checkbox, task_type, azimuth, elevation, distance,
+            focal_length, height, width, scale_factor, scale_watershed, noise_scale  # Added noise_scale here
         ],
-        outputs=[result, seed],
+        outputs=[output_images, status]
+    )
+
+    auto_caption_btn.click(
+        fn=update_prompt,
+        inputs=[images_state, task_type, molmo_message],
+        outputs=[prompt, caption_status]
+    )
+    
+    apply_mask_btn.click(
+        fn=apply_mask,
+        inputs=[images_state],
+        outputs=[output_images, status]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)