app.py

import spaces
import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor
from longcat_image.models import LongCatImageTransformer2DModel
from longcat_image.pipelines import LongCatImageEditPipeline
import numpy as np

# 1. DEFINE THE CUSTOM APPLE THEME
class AppleStyleTheme(gr.themes.Base):
    def __init__(
        self,
        # Apple Core Colors: System Blue, Cool Grays
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.gray,
        neutral_hue=gr.themes.colors.gray, 
        # Sizing: Rounded corners (iOS style) and generous spacing
        radius_size=gr.themes.sizes.radius_lg, 
        # Fonts: The stack that prioritizes San Francisco (Apple's font)
        font=[
            "system-ui", 
            "-apple-system", 
            "BlinkMacSystemFont", 
            "Segoe UI", 
            "Roboto", 
            "Helvetica Neue", 
            "Arial", 
            "sans-serif"
        ],
        font_mono=gr.themes.GoogleFont("IBM Plex Mono"),
        spacing_size=gr.themes.sizes.spacing_lg,
        text_size=gr.themes.sizes.text_md,
        **kwargs
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            radius_size=radius_size,
            font=font,
            font_mono=font_mono,
            spacing_size=spacing_size,
            text_size=text_size,
            **kwargs
        )
        
        # Override specific CSS variables for the "Apple Look"
        self.set(
            # Backgrounds: Very light gray background, pure white cards
            body_background_fill="#F5F5F7", # Apple's signature light gray
            body_text_color="#1d1d1f",      # Corrected from 'text_color'
            block_background_fill="#FFFFFF",
            block_label_background_fill="#FFFFFF",
            
            # Borders: Very subtle, thin borders
            block_border_color="#E5E5EA",
            block_border_width="1px",
            block_shadow="0 4px 6px -1px rgba(0, 0, 0, 0.05), 0 2px 4px -1px rgba(0, 0, 0, 0.03)", 
            block_padding="*spacing_lg", 
            
            # Inputs: Light gray background for inputs (like iOS grouped tables)
            input_background_fill="#F2F2F7", 
            input_border_color="transparent",
            input_shadow="none",
            input_shadow_focus="0 0 0 2px #007AFF", # Apple Blue focus ring
            
            # Buttons: Flat Apple Blue button
            button_primary_background_fill="#007AFF",
            button_primary_background_fill_hover="#0071E3",
            button_primary_text_color="white",
            button_primary_border_color="transparent",
            button_primary_shadow="none",
            
            # Text
            block_label_text_color="*neutral_500",
            block_title_text_color="*neutral_900",
        )

# Instantiate the theme
apple_theme = AppleStyleTheme()

# 2. CUSTOM CSS FOR DARK MODE FIXES
# This CSS targets the HTML elements specifically when Gradio is in "dark" mode.
custom_css = """
/* Force title color in dark mode */
.dark #app-title {
    color: #F5F5F7 !important;
}
.dark #app-subtitle {
    color: #A1A1A6 !important;
}

/* Dark mode adjustments for the custom theme colors */
.dark .gradio-container {
    background-color: #000000 !important;
}
.dark .block.panel {
    background-color: #1C1C1E !important; /* Apple Dark Mode Gray */
    border-color: #2C2C2E !important;
}
.dark input, .dark textarea, .dark .gr-input {
    background-color: #2C2C2E !important;
}
"""

# --- Model Loading ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load Image Edit Model only
edit_model_id = 'meituan-longcat/LongCat-Image-Edit'
print(f"🔄 Loading Image Edit model from {edit_model_id}...")
edit_text_processor = AutoProcessor.from_pretrained(
    edit_model_id,
    subfolder='tokenizer'
)

edit_transformer = LongCatImageTransformer2DModel.from_pretrained(
    edit_model_id,
    subfolder='transformer',
    torch_dtype=torch.bfloat16,
    use_safetensors=True
).to(device)

edit_pipe = LongCatImageEditPipeline.from_pretrained(
    edit_model_id,
    transformer=edit_transformer,
    text_processor=edit_text_processor,
)
edit_pipe.to(device, torch.bfloat16)

print(f"✅ Image Edit model loaded successfully on {device}")

# --- Core Function ---
@spaces.GPU(duration=120)
def edit_image(
    input_image: Image.Image,
    prompt: str,
    seed: int,
    progress=gr.Progress()
):
    """Edit image based on text prompt"""
    if input_image is None:
        raise gr.Error("Please upload an image first")
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter an edit instruction")
    try:
        progress(0.1, desc="Preparing image...")
        if input_image.mode != 'RGB':
            input_image = input_image.convert('RGB')
        progress(0.2, desc="Generating edited image...")
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
        with torch.inference_mode():
            output = edit_pipe(
                input_image,
                prompt,
                negative_prompt="",
                guidance_scale=4.5,
                num_inference_steps=50,
                num_images_per_prompt=1,
                generator=generator
            )
        progress(1.0, desc="Done!")
        return output.images[0]
    except Exception as e:
        raise gr.Error(f"Error during image editing: {str(e)}")

# --- Examples ---
edit_example_image_url = "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo.png"
edit_example_data = [
    [edit_example_image_url, "Add a mustache", 42],
]

# Build Gradio interface
with gr.Blocks(fill_width=True, theme=apple_theme, css=custom_css) as demo:
    
    # Header with IDs for CSS targeting
    gr.HTML("""
        <div style="text-align: center; padding: 40px 20px 30px 20px;">
            <h1 id="app-title" style="font-size: 48px; font-weight: 700; margin: 0; color: #1d1d1f; letter-spacing: -0.02em;">
                LongCat Studio
            </h1>
            <p id="app-subtitle" style="font-size: 20px; color: #86868b; margin-top: 12px; font-weight: 400;">
                AI-powered image editing
            </p>
        </div>
    """)

    with gr.Row(): 
        # Left Column: Inputs 
        with gr.Column(scale=1, variant="panel"): 
            gr.Markdown("### 🖼️ Input Image & Controls")
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                sources=["upload", "clipboard"],
                height=450,
            )

            prompt = gr.Textbox(
                label="What would you like to change?",
                placeholder="e.g., Add a mustache, Change to sunset, Make it vintage...",
                lines=2,
                max_lines=3
            )

            seed = gr.Slider(
                minimum=0,
                maximum=999999,
                value=42,
                step=1,
                label="Seed",
                visible=False
            )

            edit_btn = gr.Button("Edit Image", variant="primary", size="lg")

        # Right Column: Output 
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### ✨ Result")
            output_image = gr.Image(
                label="Result",
                type="pil",
                format="png",
                height=450,
            )

    gr.HTML("<div style='margin: 30px 0 20px 0;'></div>")

    gr.Examples(
        examples=edit_example_data,
        inputs=[input_image, prompt, seed],
        outputs=output_image,
        fn=edit_image,
        cache_examples=False,
        label="Try an example",
        examples_per_page=3
    )

    # Event handler
    edit_btn.click(
        fn=edit_image,
        inputs=[input_image, prompt, seed],
        outputs=output_image,
    )

    # Footer
    gr.HTML("""
        <div style="text-align: center; margin-top: 60px; padding: 30px 20px; border-top: 1px solid #d2d2d7;">
            <p style="color: #86868b; font-size: 13px; margin: 0;">
                Powered by LongCat • Built with 
                <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #007aff; text-decoration: none;">anycoder</a>
            </p>
        </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(mcp_server=True)