Spaces:
Running
Running
| import os | |
| import re | |
| import spaces as spaces1, spaces as spaces2 | |
| 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 as 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 | |
| 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 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 | |
| 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)) | |
| def Piper1(_do): | |
| print("starting piper1") | |
| retu = pp1( | |
| _do, | |
| height=512, | |
| width=512, | |
| num_inference_steps=4, | |
| max_sequence_length=256, | |
| guidance_scale=0 | |
| ) | |
| print("returning piper1") | |
| return retu | |
| def Piper2(img,posi,neg): | |
| retu = pp2( | |
| prompt=posi, | |
| negative_prompt=neg, | |
| image=img | |
| ) | |
| return retu | |
| 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 infer1(p): | |
| print("infer1: 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}') | |
| posi = " ".join(_do) | |
| print(posi) | |
| output = Piper1(posi) | |
| output.images[0].save(name) | |
| return name | |
| def infer2(p): | |
| print("infer2: started") | |
| p1 = p["a"] | |
| p2 = p["b"] | |
| name = p["c"] | |
| if p2 != "": | |
| _dont = f'{p2} where in {p1}' | |
| neg = _dont | |
| else: | |
| return name | |
| img = load_image(name).convert("RGB") | |
| output2 = Piper2(img,p1,neg) | |
| 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)) | |
| arr1 = [p for _ in rng] | |
| pool1 = Pool(ln) | |
| out1 = list(pool1.imap(infer1,arr1)) | |
| pool1.close() | |
| pool1.join() | |
| pool1.clear() | |
| arr2 = [{"a":p1_en,"b":p2_en,"c":out1[_]} for _ in rng] | |
| pool2 = Pool(ln) | |
| out2 = list(pool2.imap(infer2,arr2)) | |
| pool2.close() | |
| pool2.join() | |
| pool2.clear() | |
| return out2 | |
| def main(): | |
| global result | |
| global pp1 | |
| global pp2 | |
| result=[] | |
| pp1=pipe_t2i() | |
| pp2=pipe_i2i() | |
| mp.set_start_method("spawn", force=True) | |
| 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() | |
| if __name__ == "__main__": | |
| main() |