app.py

import gradio as gr
import requests
import time
import json
from contextlib import closing
from websocket import create_connection
from deep_translator import GoogleTranslator
from langdetect import detect
import os
from PIL import Image
import io
import base64
import re
from gradio_client import Client
from random import randrange
from fake_useragent import UserAgent


def flip_text(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed):
    prompt = re.sub(r'[^a-zA-Zа-яА-Я\s]', '', prompt)
    
    try:
        language = detect(prompt)
        if language == 'ru':
            prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
            print(prompt)
    except:
        pass

    try:
        with closing(create_connection(f"wss://artgan-diffusion-api.hf.space/queue/join")) as conn:
            conn.send('{"fn_index":0,"session_hash":""}')
            conn.send(f'{{"fn_index":0,"data":["{prompt}","",1024,1024,4,42,1,20,10],"session_hash":""}}')
            c = 0
            while c < 120:
                status = json.loads(conn.recv())['msg']
                if status == 'estimation':
                    time.sleep(1)
                    c += 1
                    continue
                if status == 'process_starts':
                    break            
            photo = json.loads(conn.recv())['output']['data'][0][0]['name']
            photo = "https://artgan-diffusion-api.hf.space/file=" + photo
            return photo
    except Exception as e:
        print("ERROR -->", e)
        try:
            ua = UserAgent()
            headers = {
                'authority': 'ehristoforu-stable-cascade.hf.space',
                'accept': 'text/event-stream',
                'accept-language': 'ru,en;q=0.9,la;q=0.8,ja;q=0.7',
                'cache-control': 'no-cache',
                'referer': 'https://ehristoforu-stable-cascade.hf.space/?__theme=light',
                'sec-ch-ua': '"Not_A Brand";v="8", "Chromium";v="120", "YaBrowser";v="24.1", "Yowser";v="2.5"',
                'sec-ch-ua-mobile': '?0',
                'sec-ch-ua-platform': '"Windows"',
                'sec-fetch-dest': 'empty',
                'sec-fetch-mode': 'cors',
                'sec-fetch-site': 'same-origin',
                'user-agent': f'{ua.random}'
            }
            client = Client("https://ehristoforu-stable-cascade.hf.space", headers=headers)
            result = client.predict(prompt, '', 1024, 1024, True)
            return result[0]['image']
        except:
            ua = UserAgent()
            headers = {
                'authority': 'multimodalart-stable-cascade.hf.space',
                'accept': 'text/event-stream',
                'accept-language': 'ru,en;q=0.9,la;q=0.8,ja;q=0.7',
                'cache-control': 'no-cache',
                'referer': 'https://multimodalart-stable-cascade.hf.space/?__theme=light',
                'sec-ch-ua': '"Not_A Brand";v="8", "Chromium";v="120", "YaBrowser";v="24.1", "Yowser";v="2.5"',
                'sec-ch-ua-mobile': '?0',
                'sec-ch-ua-platform': '"Windows"',
                'sec-fetch-dest': 'empty',
                'sec-fetch-mode': 'cors',
                'sec-fetch-site': 'same-origin',
                'user-agent':  f'{ua.random}'
            }
            client = Client("multimodalart/stable-cascade", headers=headers)
            result = client.predict(prompt, negative_prompt, randrange(100000), 1024, 1024, 20, 4, 10, 0, 1, api_name="/run")
            return result


def mirror(image_output, scale_by, method, gfpgan, codeformer):

    url_up = os.getenv("url_up")
    url_up_f = os.getenv("url_up_f")

    print(url_up)
    print(url_up_f)

    scale_by = int(scale_by)
    gfpgan = int(gfpgan)
    codeformer = int(codeformer)
    
    with open(image_output, "rb") as image_file:
        encoded_string2 = base64.b64encode(image_file.read())
        encoded_string2 = str(encoded_string2).replace("b'", '')

    encoded_string2 = "data:image/png;base64," + encoded_string2
    data = {"fn_index":81,"data":[0,0,encoded_string2,None,"","",True,gfpgan,codeformer,0,scale_by,512,512,None,method,"None",1,False,[],"",""],"session_hash":""}
    print(data)
    r = requests.post(f"{url_up}", json=data, timeout=100)
    print(r.text)
    ph = f"{url_up_f}" + str(r.json()['data'][0][0]['name'])
    return ph

css = """
#generate {
    width: 100%;
    background: #e253dd !important;
    border: none;
    border-radius: 50px;
    outline: none !important;
    color: white;
}
#generate:hover {
    background: #de6bda !important;
    outline: none !important;
    color: #fff;
    }
footer {visibility: hidden !important;}
#image_output {
height: 100% !important;
}
"""

with gr.Blocks(css=css) as demo:

    with gr.Tab("Базовые настройки"):
        with gr.Row():
            prompt = gr.Textbox(placeholder="Введите описание изображения...", show_label=True, label='Описание изображения:', lines=3)
    with gr.Tab("Расширенные настройки"):
        with gr.Row():
            negative_prompt = gr.Textbox(placeholder="Negative Prompt", show_label=True, label='Negative Prompt:', lines=3, value="[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry")
            seed = gr.Number(show_label=True, label="Seed:", minimum=1, maximum=1000000, value=1, step=1)
    
    with gr.Tab("Настройки апскейлинга"):
        with gr.Column():
            with gr.Row():
                scale_by = gr.Number(show_label=True, label="Во сколько раз увеличить:", minimum=1, maximum=2, value=2, step=1)
            with gr.Row():
                method = gr.Dropdown(show_label=True, value="ESRGAN_4x", label="Алгоритм увеличения", choices=["ScuNET GAN", "SwinIR 4x", "ESRGAN_4x", "R-ESRGAN 4x+", "R-ESRGAN 4x+ Anime6B"])
        with gr.Column():
            with gr.Row():
                gfpgan = gr.Slider(show_label=True, label="Эффект GFPGAN (для улучшения лица)", minimum=0, maximum=1, value=0, step=0.1)
            with gr.Row():
                codeformer = gr.Slider(show_label=True, label="Эффект CodeFormer (для улучшения лица)", minimum=0, maximum=1, value=0, step=0.1)
    
    with gr.Column():
        text_button = gr.Button("Сгенерировать изображение", variant='primary', elem_id="generate")
    with gr.Column():
        image_output = gr.Image(show_download_button=True, interactive=False, label='Результат:', elem_id='image_output', type='filepath')
        text_button.click(flip_text, inputs=[prompt, negative_prompt, seed], outputs=image_output, concurrency_limit=12)
        
        img2img_b = gr.Button("Увеличить изображение", variant='secondary')
        image_i2i = gr.Image(show_label=True, label='Увеличенное изображение:')
        img2img_b.click(mirror, inputs=[image_output, scale_by, method, gfpgan, codeformer], outputs=image_i2i)
    
#demo.launch()
demo.queue().launch(show_api=False)