Blame for app.py

import gradio as gr
import numpy as np
import argparse
import os
import random
import torch as th
th.set_float32_matmul_precision('high')
import torchvision as tv

from guided_diffusion import logger
from guided_diffusion.script_util import (
    model_and_diffusion_defaults,
    create_model_and_diffusion,
    recreate_diffusion_with_steps,
    add_dict_to_argparser,
    args_to_dict,
)
import argparse

import warnings
from torchvision import transforms
import torch.nn.functional as F
import huggingface_hub
import functools
# import spaces

# Check if running in Hugging Face Spaces
is_spaces = os.environ.get("SPACE_ID") is not None
spaces_decorator = lambda f: f  # Default no-op decorator

if is_spaces:
    try:
        import spaces
        print("Running in Hugging Face Spaces environment")
        # Create a decorator based on FREE_GPU environment variable
        free_gpu = os.environ.get("FREE_GPU", "false").lower() == "true"
        spaces_decorator = functools.partial(spaces.GPU, duration=480, free_gpu=free_gpu)
    except ImportError:
        print("spaces module not found, continuing without it")
        is_spaces = False

model_path = huggingface_hub.hf_hub_download("Kutluhan/SemanticInpaint", "model008000.pt")


def create_argparser():
    defaults = dict(
        data_dir="",
        dataset_mode="",
        clip_denoised=True,
        num_samples=10000,
        batch_size=1,
        use_ddim=False,
        model_path=model_path,
        results_path="",
        is_train=False,
        s=1.5,
        inpainting=True
    )
    defaults.update(model_and_diffusion_defaults())
    parser = argparse.ArgumentParser()
    add_dict_to_argparser(parser, defaults)
    return parser
    
args = create_argparser().parse_args()
print(args)
def load_model():
    print(args)
    args.num_classes = args.num_classes + 3 if args.inpainting else args.num_classes
    
    logger.configure()

    logger.log("creating model and diffusion...")
    model, diffusion = create_model_and_diffusion(
        **args_to_dict(args, model_and_diffusion_defaults().keys())
    )
    model.load_state_dict(
        th.load(args.model_path, map_location="cpu")
    )
    model.to("cuda")
    return model, diffusion

model, diffusion = load_model()
model.convert_to_fp16()
model.eval()

# Check if torch compile should be enabled via environment variable
use_torch_compile = os.environ.get("USE_TORCH_COMPILE", "0").lower() == "1"
if use_torch_compile:
    try:
        print("Enabling PyTorch compilation mode to improve performance")
        model = th.compile(model, mode="reduce-overhead", fullgraph=True)
    except Exception as e:
        print(f"Failed to enable PyTorch compilation mode: {e}")
        print("Continuing without compilation")

# Label to color mapping
label_color_mapping = {
    0: (0, 0, 0),           # No information
    1: (128, 64, 128),      # Urban fabric
    2: (244, 35, 232),      # Industrial / commercial
    3: (70, 70, 70),        # Mine / dump / construction
    4: (102, 102, 156),     # Artificial non-agricultural
    5: (190, 153, 153),     # Arable land
    6: (153, 153, 153),     # Permanent crops
    7: (250, 170, 30),      # Pastures
    8: (220, 220, 0),       # Complex & mixed cultivation
    9: (107, 142, 35),      # Orchards
    10: (152, 251, 152),    # Forests
    11: (70, 130, 180),     # Herbaceous vegetation
    12: (220, 20, 60),      # Open spaces
    13: (255, 0, 0),        # Wetlands
    14: (0, 0, 142),        # Water
    15: (220, 220, 220),    # Clouds & shadows
}

label_to_name = {
    0: "No information",
    1: "Urban fabric",
    2: "Industrial / commercial",
    3: "Mine / dump / construction",
    4: "Artificial non-agricultural",
    5: "Arable land",
    6: "Permanent crops",
    7: "Pastures",
    8: "Complex & mixed cultivation",
    9: "Orchards",
    10: "Forests",
    11: "Herbaceous vegetation",
    12: "Open spaces",
    13: "Wetlands",
    14: "Water",
    15: "Clouds & shadows"
}

name_to_label = {v: k for k, v in label_to_name.items()}

label_color_mapping_ts = th.tensor(list(label_color_mapping.values()))

def create_drawing_canvas():
    """Create a blank white canvas."""
    canvas = np.ones((256, 256, 3), dtype=np.uint8) * 0
    return canvas

def update_drawing_color(label):
    """
    Update the ImageEditor's brush color based on the 
    chosen label in the dropdown and show helpful hints.
    """
    # Extract label index from dropdown (e.g. "Water [USE ALONE]: (0, 0, 142)" -> 14)
    # Handle labels with tags by splitting on ":" and then extracting the base name
    label_part = label.split(":")[0].strip()
    
    # Extract the base label name without tags
    if "[USE ALONE]" in label_part:
        base_label = label_part.replace("[USE ALONE]", "").strip()
    elif "[DONT USE]" in label_part:
        base_label = label_part.replace("[DONT USE]", "").strip()
    else:
        base_label = label_part
    
    # Get the label index
    label_idx = name_to_label.get(base_label, 0)
    
    # Get RGB color from dictionary, then convert to hex (#RRGGBB)
    color = label_color_mapping[label_idx]
    hex_color = "#{:02x}{:02x}{:02x}".format(*color)
    
    # Define message types and hints for each label
    label_hints = {
        # Multi-Label Combinations That Work Well (info)
        1: ("info", "Urban Fabric: Tends to generate city-like structures and dense urban layouts. Works well with Industrial/Commercial."),
        2: ("info", "Industrial/Commercial: Often produces road networks and industrial complexes. Draw thin lines to simulate roads when using this label."),
        6: ("info", "Permanent Crops: Generates brown or green agricultural fields."),
        10: ("info", "Forests: Creates dense forested regions."),
        14: ("info", "Water: Primarily generates lakes (greenish color) and occasionally seas (bluish color). Combines well with Wetlands."),
        4: ("info", "Artificial Non-Agricultural Vegetated Areas: Complements urban layouts."),
        
        # Labels That Work Well Alone (info)
        7: ("info", "Pastures: Best used independently."),
        5: ("info", "Arable Land: Works well individually."),
        11: ("info", "Herbaceous Vegetation: Sometimes generates white snowy textures. Best used alone."),
        
        # Labels That Do Not Perform Well (warning)
        3: ("warning", "Mine/Dump/Construction: Limited generation quality due to insufficient data."),
        8: ("warning", "Complex & Mixed Cultivation Patterns: Inconsistent results."),
        9: ("warning", "Orchards: Often unreliable."),
        12: ("warning", "Open Spaces with Little or No Vegetation: Sparse or incoherent generation."),
        15: ("warning", "Clouds & Shadows: Limited performance due to insufficient data."),
        
        # Recommended Combinations (info)
        13: ("info", "Wetlands: Works best when combined with Water to generate more realistic transitions."),
        
        # Empty for other labels
        0: ("", ""),
    }
    
    # Get hint for current label
    hint_type, hint_text = label_hints.get(label_idx, ("", ""))
    
    # Return an update for the existing ImageEditor component and the hint
    return gr.ImageEditor(
        label="Semantic Drawing",
        container=True,
        brush=gr.Brush(
            colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
            default_color=hex_color,
            color_mode='fixed',
            default_size=20
        ),
        interactive=True
    ), hex_color, hint_type, hint_text

tfs_label = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.NEAREST),
    transforms.ToTensor(),
    # transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
def copy_output_to_input(output_image):
    return output_image
    
# Define the generate_image function with conditional decorator
@spaces_decorator
def generate_image(input_image, semantic_drawing, num_imgs):
    """
    Generate image using the model with adjustable prob_mask parameter.
    """    
    print("num_imgs", num_imgs)
    tfs = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    diffusion_steps = 1000
    print(diffusion_steps)
    img = tfs(input_image).unsqueeze(0)
    # Convert semantic_drawing to numpy array if it's not already
    # 256x256x3
    raw_label = semantic_drawing["composite"][:,:,:3].copy()
    # Convert RGB colors to semantic indexes
    semantic_label = np.zeros((raw_label.shape[0], raw_label.shape[1]), dtype=np.uint8)
    for label_idx, color in label_color_mapping.items():
        # Find pixels that match this color exactly
        mask = np.all(raw_label == color, axis=2)
        semantic_label[mask] = label_idx
    
    # Check if any pixels couldn't be mapped to a label
    unmapped = np.all(np.logical_not(np.any(
        [np.all(raw_label == color, axis=2) for color in label_color_mapping.values()],
        axis=0
    )), axis=0)
    
    if np.any(unmapped):
        raise ValueError(f"Found colors in the semantic drawing that don't match any label color")
    semantic_label= semantic_label.reshape(1, 1, 256, 256)
    input_label = th.FloatTensor(1, 19, 256, 256).zero_()    
    input_image = th.tensor(np.where(semantic_label == 0, img, 0))
    input_label[0, -3:, :, :] = input_image
    input_semantics = input_label.scatter_(1, th.tensor(semantic_label).long(), 1.0)
    model_kwargs = {'y': input_semantics.tile(num_imgs, 1, 1, 1),}
    model_kwargs['s'] = args.s
    print(args.s)
    
    sample_fn = (
        diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop
    )
    
    output_imgs = []    
    sample = sample_fn(
        model,
        (num_imgs, 3, 256, 256),
        clip_denoised=args.clip_denoised,
        model_kwargs=model_kwargs,
        progress=True,
    )
    sample = (sample + 1) / 2.0
    sample = sample.cpu().permute(0, 2, 3, 1)
    for num in range(num_imgs):
        output_img = sample.cpu().numpy()[num]
        output_imgs.append(output_img)
    
    return output_imgs  # return twice: for output and new input

def put_clicked_into_input(evt: gr.SelectData):
    """
    evt.value  -> value associated with the clicked item
                  (a NumPy array *or* a list, gradio versions differ)
    evt.index  -> zero-based index of the thumbnail that was clicked
    """
    clicked = evt.value
    image = clicked['image']['path']
    if isinstance(clicked, list):          # safety for older gradio versions
        clicked = clicked[evt.index]
    return image                         # single image (NumPy array)

def put_selected_into_input(chosen_img):
    print('chosen_img', chosen_img)
    if chosen_img is None:
        raise gr.Error("Please click an image in the gallery first.")
    return chosen_img
    
with gr.Blocks() as demo:
    gr.Markdown("# Satellite Image Semantic Inpainting with DDPM")
    
    # Add recommendations and best practices
    with gr.Accordion("Usage Guidelines & Best Practices", open=True):
        gr.Markdown("""
        ### General Recommendations
        
        - **Generate Multiple Samples**: Generate 2-4 samples to explore different variations. Not all outputs are visually satisfying, so multiple samples increase your chance of getting high-quality results.
        
        ### Best Practices
        
        - **Use Few Classes at a Time**: Start with 1-2 semantic classes, draw them on the mask, and generate images.
        - **Iterative Editing**: Select the best image from the generated samples as your new base, then add 1-2 more semantic classes and repeat.
        - **Preserve Context**: Using only 1-2 new mask types at a time helps maintain better visual and semantic context.
        - **Avoid Large Masks**: Drawing very large masks disrupts spatial coherence. Prefer smaller, targeted edits for better realism.
        
        ### Label Indications
        - Labels marked with **[USE ALONE]** work best independently.
        - Labels marked with **[DONT USE]** generally produce poor results.
        - Unmarked labels can be combined for interesting effects.
        """)
    
    # Create modified label names with tags
    modified_label_names = {}
    for k, v in label_to_name.items():
        if k in [7, 5, 11]:  # Labels that work well alone
            modified_label_names[k] = f"{v} [USE ALONE]"
        elif k in [3, 8, 9, 12, 15]:  # Labels that don't perform well
            modified_label_names[k] = f"{v} [DONT USE]"
        else:
            modified_label_names[k] = v
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Input Image",
                type="numpy",
                container=True,
                height=350
            )
        with gr.Column(scale=1):
            semantic_drawing = gr.ImageEditor(
                label="Semantic Drawing",
                container=True,
                brush=gr.Brush(
                    colors=[f'#{c[0]:02x}{c[1]:02x}{c[2]:02x}' for c in label_color_mapping.values()],
                    default_color='#000000',
                    color_mode='fixed',
                    default_size=20
                ),
                interactive=True,
                height=350
            )

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            label_dropdown = gr.Dropdown(
                choices=[f"{modified_label_names[k]}: {label_color_mapping[k]}" for k in label_color_mapping],
                label="Select Label",
                value=f"{modified_label_names[0]}: {label_color_mapping[0]}",
            )
        with gr.Column(scale=1):
            color_display = gr.ColorPicker(
                label="Selected Color",
                value="#000000",
                interactive=False,
                container=True
            )
    
    with gr.Row():
        num_imgs_slider = gr.Slider(
            minimum=1, maximum=4, value=1, step=1,
            label="Images to generate"
        )
    with gr.Row():
        generate_btn = gr.Button("Generate Image", size="large")
       
    
    with gr.Row():
        output_image = gr.Gallery(
            label="Generated images",
            height="auto", columns=[4],         # 4 columns looks nice up to 8 images
            preview=False,
        )

        selected_img = gr.State(value=None)
    with gr.Row():
        use_as_input_btn = gr.Button("Use Output as New Input")
        
    # Example images
    gr.Examples(
        examples=[
            ["examples/example1.png", create_drawing_canvas()],
            ["examples/example2.png", create_drawing_canvas()],
            ["examples/example3.png", create_drawing_canvas()],
            ["examples/example4.png", "./examples/example4_mask.png"],
            ["examples/example5.png", "./examples/example5_mask.png"],
            ["examples/example6.png", "./examples/example6_mask.png"],
            ["examples/example7.png", "./examples/example7_mask.png"],
            ["examples/example8.png", "./examples/example8_mask.png"],
        ],
        inputs=[input_image, semantic_drawing]
    )

    # Connect the generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[input_image, semantic_drawing, num_imgs_slider],
        outputs=output_image
    )
    
    output_image.select(
        fn=put_clicked_into_input,
        inputs=None,       # the gallery value is the *selected image*, not the list
        outputs=selected_img        # gr.Image expects a single ndarray → OK
    )

    
    use_as_input_btn.click(
        fn=put_selected_into_input,
        inputs=selected_img,     # comes from our hidden State
        outputs=input_image
    )

    
    
    # Add hint components (hidden from UI but used for handling hints)
    hint_type = gr.Textbox(visible=False)
    hint_text = gr.Textbox(visible=False)

    # Function to display appropriate notification based on hint type and text
    def show_notification(hint_type, hint_text):
        if hint_type == "warning":
            return gr.Warning(hint_text)
        elif hint_type == "info":
            return gr.Info(hint_text)
        return None

    # Connect the dropdown to update functions
    label_dropdown.change(
        fn=update_drawing_color,
        inputs=label_dropdown,
        outputs=[semantic_drawing, color_display, hint_type, hint_text]
    ).then(
        fn=show_notification,
        inputs=[hint_type, hint_text],
        outputs=None
    )

if __name__ == "__main__":
    demo.launch()