app.py

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

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
translator = Translator()
HF_TOKEN = os.environ.get("HF_TOKEN", None)
# Constants
model = "black-forest-labs/FLUX.1-dev"
default_lora = "Shakker-Labs/FilmPortrait"
default_weight_name = 'filmfotos.safetensors'
MAX_SEED = np.iinfo(np.int32).max

CSS = """
footer {
    visibility: hidden;
}
"""

JS = """function () {
  gradioURL = window.location.href
  if (!gradioURL.endsWith('?__theme=dark')) {
    window.location.replace(gradioURL + '?__theme=dark');
  }
}"""

if torch.cuda.is_available():
    device = "cuda"
    print("Using GPU")
else:
    device = "cpu"
    print("Using CPU")

pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16).to(device)
pipe.load_lora_weights(default_lora, weight_name = default_weight_name) # default load lora

def scrape_lora_link(url):
    try:
        # Send a GET request to the URL
        response = requests.get(url)
        response.raise_for_status()  # Raise an exception for bad status codes

        # Get the content of the page
        content = response.text

        # Use regular expression to find the link
        pattern = r'href="(.*?lora.*?\.safetensors\?download=true)"'
        pattern2 = r'href="(.*?\.safetensors\?download=true)"'
        match = re.search(pattern, content)
        match2 = re.search(pattern2, content)

        if match:
            safetensors_url = match.group(1)
            filename = safetensors_url.split('/')[-1].split('?')[0]  # Extract the filename from the URL
            return filename
        elif match2:
            safetensors_url = match2.group(1)
            filename = safetensors_url.split('/')[-1].split('?')[0]
            return filename
        else:
            return None
    except requests.RequestException as e:
        raise gr.Error(f"An error occurred while fetching the URL: {e}")

def enable_lora(lora_add,progress=gr.Progress(track_tqdm=True)):
    pipe.unload_lora_weights()
    if not lora_add:
        gr.Info("No Lora Loaded, Use basemodel")
        return gr.update(value="")
    else:
        url = f'https://huggingface.co/{lora_add}/tree/main'
        lora_name = scrape_lora_link(url)
        if lora_name:
            print(f'lora loading: {lora_add}/{lora_name}')
            pipe.load_lora_weights(lora_add, weight_name=lora_name)
            gr.Info(f"{lora_add} Loaded")
            return gr.update(label="LoRA Loaded Now")
        else:
            try:
                pipe.load_lora_weights(lora_add)
                print(f'lora loading: {lora_add}')
                gr.Info(f"{lora_add} Loaded")
                return gr.update(label="LoRA Loaded Now")
            except:
                raise gr.Error(f"{lora_add} Load fail, check again.")

@spaces.GPU()
def generate_image(
    prompt:str,
    lora_word:str,
    lora_scale:float=0.0,
    width:int=768,
    height:int=1024,
    scales:float=3.5,
    steps:int=24,
    seed:int=-1,
    nums:int=1,
    progress=gr.Progress(track_tqdm=True)):

    pipe.to("cuda")

    if seed == -1:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    
    prompt = str(translator.translate(prompt, 'English')) 
    text = f"{prompt} {lora_word}"

    print(f"Prompt: {text}")

    generator = torch.Generator().manual_seed(seed)
    
    image = pipe(
        prompt=text,
        height=height,
        width=width,
        guidance_scale=scales,
        output_type="pil",
        num_inference_steps=steps,
        max_sequence_length=512,
        num_images_per_prompt=nums,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images

    return image, seed
 
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. Behind,a pale gray backdrop extends,its fine texture shimmering subtly in the dim light,artfully balancing the composition and focusing attention on the subject. In a palette of black,gray,and skin tones,simplicity intertwines with profundity,as every detail whispers untold stories.",0.9,"Shakker-Labs/AWPortrait-FL",""],
    ["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. The woman has an intense and direct gaze,adding to the dramatic feel of the image. The backdrop is flowing silk,big silk. The overall composition blends elements of fashion and nature,creating a striking and powerful visual",0.9,"Shakker-Labs/AWPortrait-FL",""],
    ["A young Japanese girl, profile, blue hours, Tokyo tower",0.9,"Shakker-Labs/FilmPortrait","filmfotos, film grain, reversal film photography"],
    ["A young asian girl",0.9,"Shakker-Labs/FilmPortrait","filmfotos, film grain, reversal film photography"]
]


# Gradio Interface

with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
    gr.HTML("<h1><center>Flux Labs</center></h1>")
    gr.HTML("<p><center>Load the LoRA model on the menu</center></p>")
    with gr.Row():
        with gr.Column(scale=4):
            img = gr.Gallery(label='flux Generated Image', columns = 1, preview=True, height=600)
            with gr.Row():
                prompt = gr.Textbox(label='Enter Your Prompt (Multi-Languages)', lines=2, placeholder="Enter prompt...", scale=6)
                sendBtn = gr.Button(scale=1, variant='primary')
        with gr.Accordion("Advanced Options", open=True):
            with gr.Column(scale=1):
                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,
                )
                scales = gr.Slider(
                    label="Guidance",
                    minimum=3.5,
                    maximum=7,
                    step=0.1,
                    value=3.5,
                )
                steps = gr.Slider(
                    label="Steps",
                    minimum=1,
                    maximum=100,
                    step=1,
                    value=24,
                )
                seed = gr.Slider(
                    label="Seeds",
                    minimum=-1,
                    maximum=MAX_SEED,
                    step=1,
                    value=-1,
                )
                nums = gr.Slider(
                    label="Image Numbers",
                    minimum=1,
                    maximum=4,
                    step=1,
                    value=1,
                )
            with gr.Column(scale=1):
                lora_scale = gr.Slider(
                    label="LoRA Scale",
                    minimum=0.1,
                    maximum=1.0,
                    step=0.1,
                    value=0.9,
                )
                lora_add = gr.Textbox(
                    label="Flux LoRA",
                    info="Copy the HF LoRA model name here",
                    lines=1,
                    value="Shakker-Labs/FilmPortrait",
                )
                lora_word = gr.Textbox(
                    label="Add Flux LoRA Trigger Word",
                    info="Add the Trigger Word",
                    lines=1,
                    value="filmfotos, film grain, reversal film photography",
                )
                load_lora = gr.Button(value="Load LoRA", variant='secondary')
                
    gr.Examples(
        examples=examples,
        inputs=[prompt,lora_scale,lora_add,lora_word],
        cache_examples=False,
        examples_per_page=4,
    )

    load_lora.click(fn=enable_lora, inputs=[lora_add], outputs=lora_add)
    
    gr.on(
        triggers=[
            prompt.submit,
            sendBtn.click,
        ],
        fn=generate_image,
        inputs=[
            prompt,
            lora_word,
            lora_scale,
            width, 
            height, 
            scales, 
            steps, 
            seed, 
            nums
        ],
        outputs=[img, seed],
        api_name="run",
    )
    
demo.queue().launch()