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 LongCatImagePipeline, LongCatImageEditPipeline
import numpy as np

# 1. DEFINE THE CUSTOM THEME CLASS (Programmatic Theming, No Custom CSS)
# Inherit from Base and use .set() to achieve the Apple aesthetic.
class AppleStyleTheme(gr.themes.Base):
    def __init__(
        self,
        # Core Parameters for Apple Aesthetic
        primary_hue=gr.themes.colors.blue,
        secondary_hue=gr.themes.colors.gray,
        neutral_hue=gr.themes.colors.neutral,
        radius_size=gr.themes.sizes.radius_lg, # Large radius for modern rounded look
        font=["system-ui", "sans-serif"], # Use system fonts for native feel
        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,
            spacing_size=spacing_size,
            text_size=text_size,
            **kwargs
        )
        # Use .set() to fine-tune variables beyond the core 8
        self.set(
            # General clean white/light gray background
            body_background_fill="*neutral_50", 
            block_background_fill="white",
            # Remove default shadows for a flatter, cleaner look
            block_shadow="none",
            # Add a subtle border/ring for 'card' definition
            block_border_color="*neutral_200",
            block_border_width="1px",
            # Customize button look (Primary blue is defined by primary_hue=blue)
            button_primary_background_fill="*primary_500",
            button_primary_background_fill_hover="*primary_400",
            button_primary_border_color="*primary_500",
            # Tabs (Segmented control look is largely handled by rounded corners and soft theme)
            tabs_border_color="*neutral_200",
            tab_border_color="*neutral_100",
            tab_border_width="0",
        )

# Instantiate the custom theme
apple_theme = AppleStyleTheme()
# --- End Theme Definition ---

# --- Model Loading (Kept for completeness) ---
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 (Kept for completeness) ---
@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 (Kept for completeness) ---
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
# Passed theme=apple_theme and fill_width=True
with gr.Blocks(theme=apple_theme, fill_width=True) as demo:
    # Retaining the HTML header for branding/title
    gr.HTML("""
        <div style="text-align: center; padding: 40px 20px 30px 20px;">
            <h1 style="font-size: 48px; font-weight: 700; margin: 0; color: var(--primary-500);">
                LongCat Studio
            </h1>
            <p style="font-size: 20px; color: var(--body-text-color-subdued); margin-top: 12px; font-weight: 400;">
                AI-powered image generation and editing
            </p>
        </div>
    """)

    with gr.Tabs(selected=0):
        # Image Edit Tab (Responsive Layout: Row on Desktop, Column on Mobile)
        with gr.TabItem("Edit Image", id=0):
            with gr.Row(): 
                # Left Column (Inputs)
                with gr.Column(scale=1, min_width=0, 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,
                        info="Describe the edit you want to make to the uploaded image.",
                        max_lines=3
                    )

                    seed = gr.Slider(
                        minimum=0,
                        maximum=999999,
                        value=42,
                        step=1,
                        label="Seed",
                        visible=True # Make visible for control
                    )

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

                # Right Column (Output)
                with gr.Column(scale=1, min_width=0, variant="panel"):
                    gr.Markdown("### ✨ Result")
                    output_image = gr.Image(
                        label="Result",
                        type="pil",
                        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 (Responsive Layout: Row on Desktop, Column on Mobile)
        with gr.TabItem("Generate Image", id=1):
            with gr.Row(): 
                # Left Column (Inputs)
                with gr.Column(scale=1, min_width=0, 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
                    )

                    t2i_width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1344,
                        step=64,
                        label="Width",
                    )

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

                    t2i_seed = gr.Slider(
                        minimum=0,
                        maximum=999999,
                        value=42,
                        step=1,
                        label="Seed",
                        visible=True # Make visible for control
                    )

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

                # Right Column (Output)
                with gr.Column(scale=1, min_width=0, 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 var(--border-color-subdued);">
            <p style="color: var(--body-text-color-subdued); font-size: var(--text-sm); margin: 0;">
                Powered by LongCat • Built with 
                <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: var(--primary-500); text-decoration: none;">anycoder</a>
            </p>
        </div>
    """)

# Launch the app with the custom programmatic theme
if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=apple_theme,
    )