app.py

import gradio as gr
import io
import random
import time
import json
import base64
import requests
import os
from mistralai import Mistral
from PIL import Image
from io import BytesIO
from deep_translator import GoogleTranslator
from datetime import datetime
from theme import theme
from fastapi import FastAPI

app = FastAPI()

API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo"
API_TOKEN = os.getenv("HF_READ_TOKEN")
headers = {"Authorization": f"Bearer {API_TOKEN}"}
timeout = 100

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

def clear():
    return None 
def change_tab():
    return gr.Tabs.update(selected=1)   

# Function to query the API and return the generated image
def query(prompt, is_negative=False, steps=35, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=896, height=1152):
    if prompt == "" or prompt is None:
        return None

    key = random.randint(0, 999)
    
    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    
    # Translate the prompt from Russian to English if necessary
    prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
    print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')

    # Add some extra flair to the prompt
    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'\033[1mGeneration {key}:\033[0m {prompt}')
    
    # Prepare the payload for the API call, including width and height
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "parameters": {
            "width": width,  # Pass the width to the API
            "height": height  # Pass the height to the API
        }
    }

    # Send the request to the API and handle the response
    response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
    if response.status_code != 200:
        print(f"Error: Failed to get image. Response status: {response.status_code}")
        print(f"Response content: {response.text}")
        if response.status_code == 503:
            raise gr.Error(f"{response.status_code} : The model is being loaded")
        raise gr.Error(f"{response.status_code}")
    
    try:
        # Convert the response content into an image
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
        return image
    except Exception as e:
        print(f"Error when trying to open the image: {e}")
        return None

examples = [
    "a beautiful woman with blonde hair and blue eyes",
    "a beautiful woman with brown hair and grey eyes",
    "a beautiful woman with black hair and brown eyes",
]

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"

# CSS to style the app
css = """
footer{display:none !important}
#app-container {
    max-width: 930px;
    margin-left: auto;
    margin-right: auto;
}
"""        

# Build the Gradio UI with Blocks
with gr.Blocks(theme=theme, css=css) as app:
    # Add a title to the app
    gr.HTML("<center><h1>🎨 Stable Diffusion 3.5 large turbo 🇬🇧</h1></center>")
    with gr.Tabs() as tabs:
        with gr.TabItem(label="Image To Prompt", visible=True, id=1):
            with gr.Row():
                with gr.Column():
                    input_img = gr.Image(label="Input Picture",height=320,type="filepath")
                    submit_btn = gr.Button(value="Submit", variant='primary')
                with gr.Column():
                    output_text = gr.Textbox(label="Flux Prompt", show_copy_button = True)
                    clr_button =gr.Button("Clear",variant="primary", elem_id="clear_button")
                    clr_button.click(lambda: gr.Textbox(value=""), None, output_text)
        
        submit_btn.click(feifeichat, [input_img], [output_text])
        
        with gr.TabItem("Text to Image", visible=True, id=0):
            # Container for all the UI elements
            with gr.Column(elem_id="app-container"):
                # Add a text input for the main prompt
                with gr.Row():
                    with gr.Column(elem_id="prompt-container"):
                        with gr.Group():
                            with gr.Row():
                                text_prompt = gr.Textbox(label="Image Prompt ✍️", placeholder="Enter a prompt here", lines=2, show_copy_button = True, elem_id="prompt-text-input")
                            
                            # Accordion for advanced settings
                            with gr.Row():
                                with gr.Accordion("Advanced Settings", open=False):
                                    negative_prompt = gr.Textbox(label="Negative Prompt", lines=4, placeholder="What should not be in the image", value="(hands:-1.25), physical-defects:2, unhealthy-deformed-joints:2, unhealthy-hands:2, out of frame, bad face, (bad-image-v2-39000:1.3), (((out of frame))), deformed body features,  poor facial details, (poorly drawn face:1.3), jpeg artifacts, (missing arms:1.1), (missing legs:1.1), (extra arms:1.2), (extra legs:1.2)")
                                    with gr.Row():
                                        width = gr.Slider(label="ImageWidth", value=896, minimum=64, maximum=1216, step=32)
                                        height = gr.Slider(label="Image Height", value=1152, minimum=64, maximum=1216, step=32)
                                    steps = gr.Slider(label="Sampling steps", value=50, minimum=1, maximum=100, step=1)
                                    cfg = gr.Slider(label="CFG Scale", value=3.5, minimum=1, maximum=20, step=1)
                                    strength = gr.Slider(label="PromptStrength", value=100, minimum=0, maximum=100, step=1)
                                    seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1) # Setting the seed to -1 will make it random
                                    method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", 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"])
                # Add a button to trigger the image generation
                with gr.Row():
                    text_button = gr.Button("Generate Image 🎨", variant='primary', elem_id="gen-button")
                    clear_prompt =gr.Button("Clear Prompt 🗑️",variant="primary", elem_id="clear_button")
                    clear_prompt.click(lambda: (None), None, [text_prompt], queue=False, show_api=False)

                with gr.Group():
                    # Image output area to display the generated image
                    with gr.Row():
                        image_output = gr.Image(type="pil", label="Image Output", format="png", show_share_button=False, elem_id="gallery")
                        
                    gr.Examples(
                        examples = examples,    
                        inputs = [text_prompt],    
                    )

                with gr.Row():
                    clear_results = gr.Button(value="Clear Image 🗑️", variant="primary", elem_id="clear_button")
                    clear_results.click(lambda: (None), None, [image_output], queue=False, show_api=False)
                    
                # Bind the button to the query function with the added width and height inputs
                text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)
  

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