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
from io import BytesIO
import base64
import re
from mistralai import Mistral
from gradio_client import Client
from fake_useragent import UserAgent
import random
from theme import theme
from fastapi import FastAPI

app = FastAPI()

api_key = os.getenv("MISTRAL_API_KEY")
Mistralclient = Mistral(api_key=api_key)

def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # Open the image file
        image = Image.open(image_path).convert("RGB")

        # Resize the image to a height of 512 while maintaining the aspect ratio
        base_height = 512
        h_percent = (base_height / float(image.size[1]))
        w_size = int((float(image.size[0]) * float(h_percent)))
        image = image.resize((w_size, base_height), Image.LANCZOS)

        # Convert the image to a byte stream
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")

        return img_str
    except FileNotFoundError:
        print(f"Error: The file {image_path} was not found.")
        return None
    except Exception as e:  # Add generic exception handling
        print(f"Error: {e}")
        return None

def feifeichat(image):
    try:
        model = "pixtral-large-2411"
        # Define the messages for the chat
        base64_image = encode_image(image)
        messages = [{
            "role":
            "user",
            "content": [
                {
                    "type": "text",
                    "text": "Please provide a detailed description of this photo"
                },
                {
                    "type": "image_url",
                    "image_url": f"data:image/jpeg;base64,{base64_image}" 
                },
            ],
            "stream": False,
        }]
    
        partial_message = ""
        for chunk in Mistralclient.chat.stream(model=model, messages=messages):
            if chunk.data.choices[0].delta.content is not None:
                partial_message = partial_message + chunk.data.choices[
                    0].delta.content
                yield partial_message
    except Exception as e:  # 添加通用异常处理
        print(f"Error: {e}")
        return "Please upload a photo"


def query(prompt, negative_prompt, task, steps, sampler, cfg_scale, seed, width, height):
    result = {"prompt": prompt,"negative_prompt": negative_prompt,"task": task,"steps": steps,"sampler": sampler,"cfg_scale": cfg_scale,"seed": seed, "width": width, "height": height}
    print(result)
    try:
        language = detect(prompt)
        if language == 'ru':
            prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
            print(prompt)
    except:
        pass

    prompt = re.sub(r'[^a-zA-Zа-яА-Я\s]', '', prompt)
    
    cfg = int(cfg_scale)
    steps = int(steps)
    seed = int(seed)

    width = 896
    height = 1152

    if task == "Playground v2":
        ua = UserAgent()
        headers = {
            'user-agent': f'{ua.random}'
        }
        client = Client("https://ashrafb-arpr.hf.space/", headers=headers)
        result = client.predict(prompt, fn_index=0)
        return result
            
    if task == "Artigen v3":
        ua = UserAgent()
        headers = {
            'user-agent': f'{ua.random}'
        }
        client = Client("https://ashrafb-arv3s.hf.space/", headers=headers)
        result = client.predict(prompt,0,"Cinematic", fn_index=0)
        return result
    try:
        with closing(create_connection("wss://google-sdxl.hf.space/queue/join")) as conn:
            conn.send('{"fn_index":3,"session_hash":""}')
            conn.send(f'{{"data":["{prompt}, 4k photo","[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry",7.5,"(No style)"],"event_data":null,"fn_index":3,"session_hash":""}}')
            c = 0
            while c < 60:
                status = json.loads(conn.recv())['msg']
                if status == 'estimation':
                    c += 1
                    time.sleep(1)
                    continue
                if status == 'process_starts':
                    break
            photo = json.loads(conn.recv())['output']['data'][0][0]
            photo = photo.replace('data:image/jpeg;base64,', '').replace('data:image/png;base64,', '')
            photo = Image.open(io.BytesIO(base64.decodebytes(bytes(photo, "utf-8"))))
            return photo
    except:
        try:
            ua = UserAgent()
            headers = {
                'authority': 'ehristoforu-dalle-3-xl-lora-v2.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-dalle-3-xl-lora-v2.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("ehristoforu/dalle-3-xl-lora-v2", headers=headers)
            result = client.predict(prompt,"(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",True,0,1024,1024,6,True, api_name='/run')
            return result[0][0]['image']
        except:
            try:
                ua = UserAgent()
                headers = {
                    'authority': 'nymbo-sd-xl.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://nymbo-sd-xl.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("Nymbo/SD-XL", headers=headers)
                result = client.predict(prompt,negative_prompt,"","",True,False,False,0,896,1152,7,1,25,25,False,api_name="/run")
                return result
            except:
                try:
                    ua = UserAgent()
                    headers = {
                        'authority': 'radames-real-time-text-to-image-sdxl-lightning.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://radames-real-time-text-to-image-sdxl-lightning.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("radames/Real-Time-Text-to-Image-SDXL-Lightning", headers=headers)
                    result = client.predict(prompt, [], 0, random.randint(1, 999999), fn_index=0)
                    return result
                except:
                    try:
                        ua = UserAgent()
                        headers = {
                            'user-agent': f'{ua.random}'
                        }
                        client = Client("https://ashrafb-arpr.hf.space/", headers=headers)
                        result = client.predict(prompt, fn_index=0)
                        return result
                    except:
                        ua = UserAgent()
                        headers = {
                            'user-agent': f'{ua.random}'
                        }
                        client = Client("https://ashrafb-arv3s.hf.space/", headers=headers)
                        result = client.predict(prompt,0,"Cinematic", fn_index=0)
                        return result
                        


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

    url_up = "https://darkstorm2150-protogen-web-ui.hf.space/run/predict/"
    url_up_f = "https://darkstorm2150-protogen-web-ui.hf.space/file="

    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,896,1152,None,method,"None",1,False,[],"",""],"session_hash":""}
    r = requests.post(url_up, json=data, timeout=100)
    print(r.text)
    print(r.json()['data'][0][0]['name'])
    ph = "https://darkstorm2150-protogen-web-ui.hf.space/file=" + str(r.json()['data'][0][0]['name'])
    print(ph)
    response2 = requests.get(ph)
    img = Image.open(BytesIO(response2.content))
    return img

css = """
.gradio-container {background-color: MediumAquaMarine}
footer{display:none !important}
#generate {
    width: 100%;
}
#image_output {
height: 100% !important;
}
"""

with gr.Blocks(css=css, theme=theme) as app:
    with gr.Tab(label="Image To Prompt"):
            with gr.Row():
                with gr.Column(scale=4, min_width=300):
                    input_img = gr.Image(label="Input Picture", type="filepath")
                    
                with gr.Column(scale=3):
                    output_text = gr.Textbox(label="Flux Prompt", lines=2, scale=6, show_copy_button = True)
                    submit_btn = gr.Button(value="Generate Pompt", scale=4, variant='primary')
                    clear_prompt =gr.Button("Clear 🗑️",variant="primary", elem_id="clear_button")
                    clear_prompt.click(lambda: (None, None), None, [input_img, output_text], queue=False, show_api=False)
        
    submit_btn.click(feifeichat, [input_img], [output_text])

    with gr.Tab("Basic Settings"):
        with gr.Row():
            prompt = gr.Textbox(placeholder="Enter the image description...", show_label=True, label='Image Prompt ✍️', lines=3, show_copy_button = True)
        with gr.Row():
            task = gr.Radio(interactive=True, value="Stable Diffusion XL 1.0", show_label=True, label="Model of neural network ❄️ ", choices=['Stable Diffusion XL 1.0', 'Crystal Clear XL', 
                                                                                                              'Juggernaut XL', 'DreamShaper XL',
                                                                                                              'SDXL Niji', 'Cinemax SDXL', 'NightVision XL'])
    #with gr.Tab("Extended settings"):
        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")
        with gr.Row():
            sampler = gr.Dropdown(value="DPM++ 2M Karras", show_label=True, label="Sampling Method:", choices=[
                "DPM++ 2M Karras", "DPM++ 2S a Karras", "DPM2 a Karras", "DPM2 Karras", "DPM++ SDE Karras", "DEIS", "LMS", "DPM Adaptive", "DPM++ 2M", "DPM2 Ancestral", "DPM++ S", "DPM++ SDE", "DDPM", "DPM Fast", "dpmpp_2s_ancestral", "Euler", "Euler CFG PP", "Euler a", "Euler Ancestral", "Euler+beta", "Heun", "Heun PP2", "DDIM", "LMS Karras", "PLMS", "UniPC", "UniPC BH2"])
        with gr.Row():
            steps = gr.Slider(show_label=True, label="Sampling Steps:", minimum=1, maximum=50, value=35, step=1)
        with gr.Row():
            cfg_scale = gr.Slider(show_label=True, label="CFG Scale:", minimum=1, maximum=20, value=7, step=1)
        with gr.Row():
            seed = gr.Number(show_label=True, label="Seed:", minimum=-1, maximum=1000000, value=-1, step=1)
        with gr.Row():
            width = gr.Slider(label="Width", minimum=512, maximum=2048, step=8, value=896, interactive=True,)
        with gr.Row():
            height = gr.Slider(label="Height", minimum=512, maximum=2048, step=8, value=1152,interactive=True,)
            with gr.Column():
                text_button = gr.Button("Generate image", variant="primary", interactive=True, elem_id="generate")
                image_output = gr.Image(show_download_button=True, interactive=False, label='Generated Image 🌄', show_share_button=False, format="png", type="filepath", elem_id="gallery")
            
    with gr.Tab("Upscaling Settings"):
        with gr.Column():
            with gr.Row():
                scale_by = gr.Number(show_label=True, label="How many times to increase:", minimum=1, maximum=2, value=2, step=1)
            with gr.Row():
                method = gr.Dropdown(show_label=True, value="ESRGAN_4x", label="Increasing algorithm", 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="Effect GFPGAN (For facial improvement)", minimum=0, maximum=1, value=0, step=0.1)
            with gr.Row():
                codeformer = gr.Slider(show_label=True, label="Effect CodeFormer (Improve the face)", minimum=0, maximum=1, value=0, step=0.1)
    

    with gr.Column():
        #image_output = gr.Image(show_download_button=True, interactive=False, label='Generated Image 🌄', show_share_button=False, format="png", type="filepath", elem_id="gallery")
        
        text_button.click(query, inputs=[prompt, negative_prompt, task, steps, sampler, cfg_scale, seed, width, height], outputs=image_output, concurrency_limit=48)
        clear_prompt =gr.Button("Clear 🗑️",variant="primary", elem_id="clear_button")
        clear_prompt.click(lambda: (None, None), None, [prompt, image_output], queue=False, show_api=False)
        
        img2img_b = gr.Button("Increase the image", variant='secondary')
        image_i2i = gr.Image(show_label=True, label='Increased image:')
        img2img_b.click(mirror, inputs=[image_output, scale_by, method, gfpgan, codeformer], outputs=image_i2i, concurrency_limit=48)
        

app.queue(default_concurrency_limit=200, max_size=200)  # <-- Sets up a queue with default parameters
if __name__ == "__main__":
    app.launch()