app.py

import os
import re
import spaces
import random
import string
import torch
import requests
import gradio as gr
import numpy as np
from lxml.html import fromstring
#from transformers import pipeline
from torch import multiprocessing as mp
from torch.multiprocessing import Pool
#from pathos.multiprocessing import ProcessPool as Pool
#from pathos.threading import ThreadPool as Pool
#from diffusers.pipelines.flux import FluxPipeline
#from diffusers.utils import export_to_gif
#from huggingface_hub import hf_hub_download
#from safetensors.torch import load_file
from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
from diffusers.utils import load_image
#import jax
#import jax.numpy as jnp

result=[]

def port_inc():
    env = os.getenv("CUSTOM_PORT")
    if env == None:
        os.environ["CUSTOM_PORT"]="7860"
    else:
        os.environ["CUSTOM_PORT"]=str(int(env)+1)

def init_pool():
    #port_inc()
    print("init pool")

def pipe_t2i():
    PIPE = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, token=os.getenv("hf_token")).to("cuda")
    return PIPE
    
def pipe_i2i():
    PIPE = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
    PIPE.unet = torch.compile(PIPE.unet, mode="reduce-overhead", fullgraph=True)
    return PIPE

pp1=pipe_t2i()
pp2=pipe_i2i()

def translate(text,lang):
    if text == None or lang == None:
        return ""       
    text = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', text)).lower().strip()
    lang = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', lang)).lower().strip()    
    if text == "" or lang == "":
        return ""
    if len(text) > 38:
        raise Exception("Translation Error: Too long text!")
    user_agents = [
        'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15',
        'Mozilla/5.0 (Macintosh; Intel Mac OS X 13_1) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.1 Safari/605.1.15'
    ]
    padded_chars = re.sub("[(^\-)(\-$)]","",text.replace("","-").replace("- -"," ")).strip()
    query_text = f'Please translate {padded_chars}, into {lang}'
    url = f'https://www.google.com/search?q={query_text}'
    resp = requests.get(
        url = url,
        headers = {
            'User-Agent': random.choice(user_agents)
        }
    )
    content = resp.content
    html = fromstring(content)
    translated = text
    try:
        src_lang = html.xpath('//*[@class="source-language"]')[0].text_content().lower().strip()
        trgt_lang = html.xpath('//*[@class="target-language"]')[0].text_content().lower().strip()
        src_text = html.xpath('//*[@id="tw-source-text"]/*')[0].text_content().lower().strip()
        trgt_text = html.xpath('//*[@id="tw-target-text"]/*')[0].text_content().lower().strip()
        if trgt_lang == lang:
            translated = trgt_text
    except:
        print(f'Translation Warning: Failed To Translate!')
    ret = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', translated)).lower().strip()
    print(ret)
    return ret

def generate_random_string(length):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))

@spaces.GPU(duration=40)
def Piper1(_do):
    try:
        retu = pp1(
            _do,
            height=512,
            width=512,
            num_inference_steps=4,
            max_sequence_length=256,
            guidance_scale=0
        )
        return retu
    except Exception as e:
        print(e)
        return None

@spaces.GPU(duration=40)
def Piper2(img,posi,neg):
    try:
        retu = pp2(
            prompt=posi,
            negative_prompt=neg,
            image=img
        )
        return retu
    except Exception as e:
        print(e)
        return None

@spaces.GPU(duration=40)
def tok(pipe,txt):
    toks = pipe.tokenizer(txt)['input_ids']
    print(toks)
    return toks

css="""
input, input::placeholder {
    text-align: center !important;
}
*, *::placeholder {
    direction: ltr !important;
    font-family: Suez One !important;
}
h1,h2,h3,h4,h5,h6,span,p,pre {
    width: 100% !important;
    text-align: center !important;
    display: block !important;
}
footer {
    display: none !important;
}
#col-container {
    margin: 0 auto !important;
    max-width: 15cm !important;
 }
.image-container {
    aspect-ratio: 512 / 512 !important;
}
.dropdown-arrow {
    display: none !important;
}
*:has(.btn), .btn {
    width: 100% !important;
    margin: 0 auto !important;
}
"""

js="""
function custom(){
    document.querySelector("div#prompt input").setAttribute("maxlength","38")
    document.querySelector("div#prompt2 input").setAttribute("maxlength","38")
}
"""

def infer(p):
        print("infer: started")
        p1 = p["a"]
        print(f'prompt 1: {p1}')
        p2 = p["b"]
        print(f'prompt 2: {p2}')
        name = generate_random_string(12)+".png"
        print(f'name: {name}')
        _do = ['beautiful', 'playful', 'photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable coloring', 'rough texture', 'best quality', 'focused']
        if p1 != "":
            _do.append(f'{p1}')
        if p2 != "":
            _dont = f'{p2} where in {p1}'
            neg = _dont
        else:
            neg = None
        posi = " ".join(_do)
        print(posi)
        output = Piper1(posi)
        print(output)
        if output == None:
            return None
        else:
            output.images[0].save(name)
        if neg == None:
            return name
        img = load_image(name).convert("RGB")
        output2 = Piper2(img,posi,neg)
        if output2 == None:
            return None
        else:
            output2.images[0].save("_"+name)
        return "_"+name

def run(p1,p2,*result):
        p1_en = translate(p1,"english")
        p2_en = translate(p2,"english")
        p = {"a":p1_en,"b":p2_en}
        ln = len(result)
        print("images: "+str(ln))
        rng = list(range(ln))
        arr = [p for _ in rng]
        out = None
        print("Starting pool!")
        with Pool(ln, initializer=init_pool) as pool:
            out = pool.imap(infer,arr)
        return list(out)

def ui():
    with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
        with gr.Column(elem_id="col-container"):
            gr.Markdown(f"""
                # MULTI-LANGUAGE IMAGE GENERATOR
            """)
            with gr.Row():
                prompt = gr.Textbox(
                    elem_id="prompt",
                    placeholder="INCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                prompt2 = gr.Textbox(
                    elem_id="prompt2",
                    placeholder="EXCLUDE",
                    container=False,
                    max_lines=1
                )
            with gr.Row():
                run_button = gr.Button("START",elem_classes="btn",scale=0)
            with gr.Row():
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
                result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))

        run_button.click(fn=run,inputs=[prompt,prompt2,*result],outputs=result)
        demo.queue().launch(server_port=int(os.getenv("CUSTOM_PORT")))

if __name__ == "__main__":

    mp.set_start_method("fork", force=True)
    port_inc()
    ui()