app.py

import spaces
import gradio as gr
import torch
from diffusers import FluxPipeline
from huggingface_hub import hf_hub_download
from PIL import Image
import requests
from translatepy import Translator
import numpy as np
import random

# Initialize Translator
translator = Translator()

# Constants
model = "Shakker-Labs/AWPortrait-FL"
MAX_SEED = np.iinfo(np.int32).max

# Device & dtype setup
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

# Load model and move to device
pipe = FluxPipeline.from_pretrained(model, torch_dtype=dtype)
pipe.to(device)

# Image generation function
@spaces.GPU()
def generate_image(
    prompt,
    width=768,
    height=1024,
    scale=3.5,
    steps=24,
    seed=-1,
    nums=1,
    progress=gr.Progress(track_tqdm=True)
):
    if seed == -1:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

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

    prompt = str(translator.translate(prompt, 'English'))
    print(f'Prompt: {prompt}')
    
    image = pipe(
        prompt, 
        width=width,
        height=height,
        guidance_scale=scale,
        num_inference_steps=steps,
        generator=generator,
        output_type="pil",
        max_sequence_length=512,
        num_images_per_prompt=nums,
    ).images

    return image, seed

# Example prompts
examples = [
    "close up portrait, Amidst the interplay of light and shadows in a photography studio,a soft spotlight traces the contours of a face,highlighting a figure clad in a sleek black turtleneck. The garment,hugging the skin with subtle luxury,complements the Caucasian model's understated makeup,embodying minimalist elegance.",
    "Caucasian,The image features a young woman of European descent standing in an studio setting,surrounded by silk. (She is wearing a silk dress),paired with a bold. Her brown hair is wet and tousled,falling naturally around her face,giving her a raw and edgy look.",
    "A black and white portrait of a young woman with a captivating gaze. She's bundled up in a cozy black sweater,hands gently cupped near her face.",
    "Fashion photography portrait,close up portrait,(a woman of European descent is surrounded by lava rock and magma from head to neck, red magma hair, wear volcanic lava rock magma outfit coat lava rock magma fashion costume with ruffled layers"
]

# Gradio Interface
css = """
footer {
    visibility: hidden;
}
"""

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:

    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(label='Enter Your Prompt (multilingual)', scale=6)
            submit = gr.Button(scale=1, variant='primary')

    img = gr.Gallery(label="Gallery", columns=1, preview=True)

    with gr.Accordion("Advanced Options", open=False):
        with gr.Row():
            width = gr.Slider(label="Width", minimum=512, maximum=1280, step=8, value=768)
            height = gr.Slider(label="Height", minimum=512, maximum=1280, step=8, value=1024)
        with gr.Row():
            scale = gr.Slider(label="Guidance Scale", minimum=0, maximum=50, step=0.1, value=3.5)
            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=24)
        with gr.Row():
            seed = gr.Slider(label="Seed (-1 = Random)", minimum=-1, maximum=MAX_SEED, step=1, value=0, visible=True)
            nums = gr.Slider(label="Image Numbers", minimum=1, maximum=4, step=1, value=1, scale=1)

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[img, seed],
        fn=generate_image,
        cache_examples="lazy",
    )

    gr.on(
        triggers=[prompt.submit, submit.click],
        fn=generate_image,
        inputs=[prompt, width, height, scale, steps, seed, nums],
        outputs=[img, seed],
        api_name="run",
    )

demo.queue().launch()