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, LongCatImagePipeline
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')

# Text-to-Image Model
t2i_model_id = 'meituan-longcat/LongCat-Image'
print(f"🔄 Loading Text-to-Image model from {t2i_model_id}...")
t2i_text_processor = AutoProcessor.from_pretrained(
    t2i_model_id,
    subfolder='tokenizer'
)

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

t2i_pipe = LongCatImagePipeline.from_pretrained(
    t2i_model_id,
    transformer=t2i_transformer,
    text_processor=t2i_text_processor,
)
t2i_pipe.to(device, torch.bfloat16)

print(f"✅ Text-to-Image model loaded successfully")

# Image Edit Model
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 Functions ---
@spaces.GPU(duration=120)
def generate_image(
    prompt: str,
    width: int,
    height: int,
    seed: int,
    progress=gr.Progress()
):
    """Generate image from text prompt"""
    if not prompt or prompt.strip() == "":
        raise gr.Error("Please enter a prompt")
    try:
        progress(0.1, desc="Preparing generation...")
        progress(0.2, desc="Generating image...")
        generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
        with torch.inference_mode():
            output = t2i_pipe(
                prompt,
                negative_prompt="",
                height=height,
                width=width,
                guidance_scale=4.5,
                num_inference_steps=50,
                num_images_per_prompt=1,
                generator=generator,
                enable_cfg_renorm=True,
                enable_prompt_rewrite=True
            )
        progress(1.0, desc="Done!")
        return output.images[0]
    except Exception as e:
        raise gr.Error(f"Error during image generation: {str(e)}")

@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],
]

t2i_example_prompts = [
    ["一个年轻的亚裔女性，身穿黄色针织衫，搭配白色项链。她的双手放在膝盖上，表情恬静。背景是一堵粗糙的砖墙，午后的阳光温暖地洒在她身上，营造出一种宁静而温馨的氛围。", 1344, 768, 43],
    ["A serene mountain landscape at sunset with golden clouds", 1344, 768, 42],
    ["A cute robot sitting at a desk, digital art style", 1024, 1024, 44],
]

# Build Gradio interface
# Theme and CSS are PASSED IN LAUNCH()
with gr.Blocks(fill_width=True) 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 generation and editing
            </p>
        </div>
    """)

    with gr.Tabs(selected=0):
        # Image Edit Tab
        with gr.TabItem("Edit Image", id=0):
            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
            )

        # Text-to-Image Tab
        with gr.TabItem("Generate Image", id=1):
            with gr.Row(): 
                # Left Column: Inputs
                with gr.Column(scale=1, variant="panel"):
                    gr.Markdown("### 🎨 Generation Controls")
                    t2i_prompt = gr.Textbox(
                        label="Describe your image",
                        placeholder="e.g., A serene mountain landscape at sunset...",
                        lines=4,
                        max_lines=6
                    )

                    # Hidden Sliders
                    t2i_width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1344,
                        step=64,
                        label="Width",
                        visible=False,
                    )

                    t2i_height = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=768,
                        step=64,
                        label="Height",
                        visible=False,
                    )

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

                    generate_btn = gr.Button("Generate Image", variant="primary", size="lg")

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

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

            gr.Examples(
                examples=t2i_example_prompts,
                inputs=[t2i_prompt, t2i_width, t2i_height, t2i_seed],
                outputs=t2i_output,
                fn=generate_image,
                cache_examples=False,
                label="Try an example",
                examples_per_page=3
            )

    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[t2i_prompt, t2i_width, t2i_height, t2i_seed],
        outputs=t2i_output,
    )

    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
# THEME and CSS moved here for Gradio 6 Compliance
if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=apple_theme,
        css=custom_css
    )