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 diffusers import AutoPipelineForText2Image
#from tempfile import NamedTemporaryFile
#from pathos.multiprocessing import ProcessingPool as ProcessPoolExecutor

#pool = ProcessPoolExecutor(1000)
#pool.__enter__()

#model_id = "runwayml/stable-diffusion-v1-5"
#model_id = "kandinsky-community/kandinsky-3"
model_id = "stabilityai/stable-diffusion-3-medium-diffusers"

device = "cuda" if torch.cuda.is_available() else "cpu"
if torch.cuda.is_available():
    torch.cuda.max_memory_allocated(device=device)
    pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True, token=os.getenv('hf_token'))
    pipe = pipe.to(device)
else: 
    pipe = AutoPipelineForText2Image.from_pretrained(model_id, use_safetensors=True, token=os.getenv('hf_token'))
    pipe = pipe.to(device)

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}'

    print(url)
    
    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=45)
def Piper(_do,_dont):
    return pipe(
        _do,
        height=512,
        width=576,
        negative_prompt=_dont,
        num_inference_steps=200,
        guidance_scale=4
    )

def infer(prompt,prompt2):
    name = generate_random_string(12)+".png"
    prompt_en = translate(prompt,"english")
    prompt2_en = translate(prompt2,"english")
    if prompt == None or prompt.strip() == "":
        _do = 'warm dark muted vintage foreground colors, rough texture, dynamic poze, proportional, reasonable, realistic, natural, cold light vivid playful background colors'
    else:
        _do = f'photographed { prompt_en }, warm dark muted vintage foreground colors, rough texture, dynamic poze, proportional, reasonable, realistic, natural, cold light vivid playful background colors'
    if prompt2 == None or prompt2.strip() == "":
        _dont = 'unreasonable, unreal, ugly, deformed, disfigured, poor details, bad anatomy, logos, texts, labels'
    else:
        _dont = f'{prompt2_en}, {prompt2_en} where in {prompt_en}, labels where in {prompt_en}, unreasonable, unreal, texts, logos, ugly, deformed, disfigured, poor details, bad anatomy'
    image = Piper(_do,_dont).images[0].save(name)
    return name

css="""
input::placeholder {
    text-align: center !important;
}
*, *::placeholder {
    direction: rtl !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: 576 / 512 !important;
}
.dropdown-arrow {
    display: none !important;
}
"""

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

if torch.cuda.is_available():
    power_device = "מעבד גרפי"
else:
    power_device = "מעבד לוגי"

with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
            # תמונות איכותיות מבינה מלאכותית
            רץ על {power_device}
        """)
        with gr.Row():
            prompt = gr.Textbox(
                elem_id="prompt",
                placeholder="מה *כן* להוסיף",
                container=False,
                rtl=True,
                max_lines=1
            )
        with gr.Row():
            prompt2 = gr.Textbox(
                elem_id="prompt2",
                placeholder="מה *לא* להוסיף",
                container=False,
                rtl=True,
                max_lines=1
            )
        with gr.Row():
            run_button = gr.Button("התחלה")
        result = gr.Image(elem_id="image-container", label="Result", show_label=False, type='filepath', show_share_button=False)
    prompt.submit(
        fn = infer,
        inputs = [prompt,prompt2],
        outputs = [result]
    )
    prompt2.submit(
        fn = infer,
        inputs = [prompt,prompt2],
        outputs = [result]
    )
    run_button.click(
        fn = infer,
        inputs = [prompt,prompt2],
        outputs = [result]
    )

demo.queue().launch()