app.py

import gradio as gr
import numpy as np
import random
import json

import spaces
import torch
from pikigen import PikigenPipeline

# Trick required because it is not a native diffusers model
from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
LOADABLE_CLASSES["pikigen"] = LOADABLE_CLASSES["pikigen.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]

device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "Freepik/Pikigen-test"

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float32

pipe = PikigenPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
# pipe.enable_model_cpu_offload()  # For less memory consumption
pipe.to(device)
pipe.vae.enable_slicing()
pipe.vae.enable_tiling()

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1600

# Predefined resolutions
RESOLUTIONS = {
  "horizontal": [
    {"width": 1344, "height": 896, "label": "1344×896"},
    {"width": 1152, "height": 768, "label": "1152×768"},
    {"width": 960, "height": 640, "label": "960×640"},
    {"width": 1600, "height": 896, "label": "1600×896"}
  ],
  "vertical": [
    {"width": 896, "height": 1344, "label": "896×1344"},
    {"width": 768, "height": 1152, "label": "768×1152"},
    {"width": 640, "height": 960, "label": "640×960"},
    {"width": 896, "height": 1600, "label": "896×1600"}
  ],
  "square": [
    {"width": 1216, "height": 1216, "label": "1216×1216"},
    {"width": 1024, "height": 1024, "label": "1024×1024"}
  ],
  "default": {"width": 1024, "height": 1024, "label": "1024×1024"}
}

# Create flattened options for the dropdown
resolution_options = []
for category, resolutions in RESOLUTIONS.items():
    if category != "default":
        for res in resolutions:
            resolution_options.append(f"{category.capitalize()} - {res['label']}")


@spaces.GPU(duration=120)
def infer(
    prompt,
    negative_prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator().manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image, seed


def update_resolution(resolution_option):
    """Updates width and height based on selected resolution option"""
    if not resolution_option:
        # Use default resolution
        return RESOLUTIONS["default"]["width"], RESOLUTIONS["default"]["height"]
    
    # Parse the resolution option format: "Category - WidthxHeight"
    try:
        category, label = resolution_option.split(" - ")
        category = category.lower()
        
        for res in RESOLUTIONS[category]:
            if res["label"] == label:
                return res["width"], res["height"]
    except:
        pass
    
    # Fallback to default
    return RESOLUTIONS["default"]["width"], RESOLUTIONS["default"]["height"]


examples = [
    "A photorealistic 3D render of a charming, mischievous young boy, approximately eight years old, possessing the endearingly unusual features of long, floppy donkey ears that droop playfully over his shoulders and a surprisingly small, pink pig nose that twitches slightly.  His eyes, a sparkling, intelligent hazel, are wide with a hint of playful mischief, framed by slightly unruly, sandy-brown hair that falls in tousled waves across his forehead.  He's dressed in a simple, slightly oversized, worn denim shirt and patched-up corduroy trousers, hinting at a life spent playing outdoors. The lighting is soft and natural, casting gentle shadows that highlight the texture of his skin – slightly freckled and sun-kissed, suggesting time spent in the sun.  His expression is one of curious anticipation, his lips slightly parted as if he’s about to speak or perhaps is listening intently. The background is a subtly blurred pastoral scene, perhaps a sun-dappled meadow with wildflowers, enhancing the overall whimsical and slightly surreal nature of the character.  The overall style aims for a blend of realistic rendering with a touch of whimsical cartoonishness, capturing the unique juxtaposition of the boy's human features and his animalistic ears and nose.",
    "Two white swans with long necks, gracefully swimming in a still body of water. The swans are positioned in a heart shape, with their necks intertwined, creating a romantic and elegant scene. The water is calm and reflective, reflecting the soft, golden light of the setting sun. The background is a blur of soft, golden hues, suggesting a peaceful and serene environment. The image is likely a photograph, captured with a shallow depth of field, which emphasizes the swans and creates a sense of intimacy. The soft lighting and the gentle curves of the swans create a sense of tranquility and beauty. The overall mood of the image is one of love, peace, and serenity.",
    """A watercolor painting of the American flag waving in the wind. The flag is painted in a vibrant red, white, and blue, with the stars in the blue field appearing slightly blurred, creating a sense of motion. The red stripes are painted with a slightly textured brushstroke, giving the flag a realistic and weathered look. The flag is positioned diagonally across the image, with the top left corner extending beyond the frame. The background is a simple white, allowing the flag to be the focal point. Below the flag, in bold red letters, is the text "PRESIDENTS DAY," with "21 FEBRUARY" in blue text above it. Below the text, in black, is "UNITED STATES OF AMERICA." The overall style of the image is patriotic and celebratory, with the watercolor technique adding a touch of artistic flair. The image evokes a sense of pride and national unity, making it a fitting tribute to Presidents Day.""",
    "A captivating photo, shot with a shallow depth of field, of a stunning blonde woman with cascading waves of platinum blonde hair that fall past her shoulders, catching the light. Her eyes, a striking shade of emerald green, are intensely focused on something just off-camera, creating a sense of intrigue.  Sunlight streams softly onto her face, highlighting the delicate curve of her cheekbones and the subtle freckles scattered across her nose. She's wearing a flowing, bohemian-style maxi dress, the fabric a deep sapphire blue that complements her hair and eyes beautifully. The dress is adorned with intricate embroidery along the neckline and sleeves, adding a touch of elegance.  The background is intentionally blurred, suggesting a sun-drenched garden setting with hints of vibrant flowers and lush greenery, drawing the viewer's eye to the woman's captivating features.  The overall mood is serene yet captivating, evoking a feeling of summer warmth and quiet contemplation.  The image should have a natural, slightly ethereal quality, with soft, diffused lighting that enhances her beauty without harsh shadows.",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(" # F-lite Text-to-Image Demo")

        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )

            run_button = gr.Button("Run", scale=0, variant="primary")

        result = gr.Image(label="Result", show_label=False)

        with gr.Accordion("Advanced Settings", open=False):
            with gr.Tabs() as resolution_tabs:
                with gr.TabItem("Preset Resolutions"):
                    resolution_dropdown = gr.Dropdown(
                        label="Select Resolution",
                        choices=[""] + resolution_options,
                        value="",
                        allow_custom_value=False,
                    )
                
                with gr.TabItem("Custom Resolution"):
                    with gr.Row():
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=RESOLUTIONS["default"]["width"],
                        )

                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=RESOLUTIONS["default"]["height"],
                        )

            negative_prompt = gr.Text(
                label="Negative prompt",
                max_lines=1,
                placeholder="Enter a negative prompt",
                visible=True,
            )

            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )

            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance scale",
                    minimum=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=3.5,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=30,
                )

        # Examples should explicitly target only the prompt input
        gr.Examples(examples=examples, inputs=prompt, example_labels=[ex[:120] + "..." if len(ex) > 120 else ex for ex in examples])
    
    # Update width and height when resolution is selected from dropdown
    resolution_dropdown.change(
        fn=update_resolution,
        inputs=[resolution_dropdown],
        outputs=[width, height]
    )
    
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

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