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

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

# --- Custom CSS (Modified to ensure Row behavior) ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=SF+Pro+Display:wght@300;400;500;600;700&display=swap');

* {
    font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
}

/* Ensure padding on all screen sizes */
.gradio-container {
    /* Retain max-width but remove side padding in Blocks for better full-width control */
    max-width: 1400px !important;
    margin: auto !important;
    padding: 0 16px; 
}

/* CSS Fix: Ensure layout elements behave as flex containers */
.gr-row {
    display: flex !important;
}

/* Background gradient for the overall app (like a subtle card) */
#component-0 {
    background: linear-gradient(180deg, #f5f5f7 0%, #ffffff 100%) !important;
}

/* Tab bar styling for the segmented control look */
.tabs {
    border: none !important;
    background: transparent !important;
}

.tab-nav {
    border: none !important;
    background: rgba(255, 255, 255, 0.8) !important;
    backdrop-filter: blur(20px) !important;
    border-radius: 12px !important;
    padding: 4px !important;
    gap: 4px !important;
}

button.selected {
    background: white !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08) !important;
    border-radius: 8px !important;
    font-weight: 500 !important;
}

/* Image and input component styling */
.input-image, .output-image {
    border-radius: 16px !important;
    overflow: hidden !important;
    box-shadow: 0 4px 16px rgba(0, 0, 0, 0.06) !important;
}

textarea, input[type="text"] {
    border: 1px solid #d2d2d7 !important;
    border-radius: 12px !important;
    padding: 12px 16px !important;
    font-size: 15px !important;
    transition: all 0.2s ease !important;
}

textarea:focus, input[type="text"]:focus {
    border-color: #007aff !important;
    box-shadow: 0 0 0 3px rgba(0, 122, 255, 0.1) !important;
}

/* Primary Button Styling */
.primary-btn {
    background: linear-gradient(180deg, #007aff 0%, #0051d5 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 14px 28px !important;
    font-size: 16px !important;
    font-weight: 500 !important;
    color: white !important;
    box-shadow: 0 4px 12px rgba(0, 122, 255, 0.3) !important;
    transition: all 0.2s ease !important;
}

.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 16px rgba(0, 122, 255, 0.4) !important;
}

/* Slider and Label styling */
.slider-container {
    margin: 16px 0 !important;
}

label {
    font-weight: 500 !important;
    color: #1d1d1f !important;
    margin-bottom: 8px !important;
    font-size: 14px !important;
}

.secondary-text {
    color: #86868b !important;
    font-size: 13px !important;
}

/* Card Style (targets gr-panel when variant="panel" is used) */
.card {
    background: white !important;
    border-radius: 16px !important;
    padding: 24px !important;
    box-shadow: 0 2px 12px rgba(0, 0, 0, 0.04) !important;
}

/* Mobile adjustments */
@media (max-width: 768px) {
    .gradio-container {
        padding: 0 8px !important;
    }
    .card {
        padding: 16px !important;
    }
}
"""

# Build Gradio interface
# Added fill_width=True to Blocks to maximize space utilization
with gr.Blocks(fill_width=True) as demo:
    gr.HTML("""
        <div style="text-align: center; padding: 40px 20px 30px 20px;">
            <h1 style="font-size: 48px; font-weight: 700; margin: 0; background: linear-gradient(90deg, #007aff 0%, #5856d6 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">
                LongCat Studio
            </h1>
            <p 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 (Responsive Layout: Row on Desktop, Column on Mobile)
        with gr.TabItem("Edit Image", id=0):
            with gr.Row(): 
                # Left Column (Inputs)
                # Setting min_width=0 ensures the column will shrink freely on narrow viewports
                # until Gradio's responsive behavior naturally stacks them.
                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,
                        elem_classes=["input-image"]
                    )

                    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", elem_classes=["primary-btn"])

                # 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,
                        elem_classes=["output-image"]
                    )

            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=False
                    )

                    generate_btn = gr.Button("Generate Image", variant="primary", size="lg", elem_classes=["primary-btn"])

                # 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,
                        elem_classes=["output-image"]
                    )

            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 with theme and custom CSS
if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        theme=gr.themes.Soft(), 
        css=custom_css 
    )