app.py

from PIL import Image
import base64
from io import BytesIO
import os
from mistralai import Mistral
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from huggingface_hub import hf_hub_download
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from openai import OpenAI
import config
from extras.expansion import FooocusExpansion
import re

expansion = FooocusExpansion()

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

client_open_ai = OpenAI(
    base_url="https://api-inference.huggingface.co/v1/",
    api_key=os.getenv('HF_TOKEN')
)

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)

good_vae = AutoencoderKL.from_pretrained("shuttleai/shuttle-3-diffusion", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("shuttleai/shuttle-3-diffusion", torch_dtype=dtype, vae=taef1).to(device)

pipe.load_lora_weights(hf_hub_download("aifeifei798/feifei-flux-lora-v1", "feifei.safetensors"), adapter_name = "feifei")
pipe.load_lora_weights(hf_hub_download("aifeifei798/feifei-flux-lora-v1", "FLUX-dev-lora-add_details.safetensors"), adapter_name = "FLUX-dev-lora-add_details")
pipe.load_lora_weights(hf_hub_download("aifeifei798/feifei-flux-lora-v1", "Shadow-Projection.safetensors"), adapter_name = "Shadow-Projection")
pipe.set_adapters(["feifei","FLUX-dev-lora-add_details","Shadow-Projection"], adapter_weights=[0.65,0.35,0.35])
pipe.fuse_lora(adapter_name=["feifei","FLUX-dev-lora-add_details","Shadow-Projection"], lora_scale=1.0)

pipe.unload_lora_weights()
torch.cuda.empty_cache()


MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 4096

css="""
#col-container {
    width: auto;
    height: 750px;
}
"""
@spaces.GPU()
def infer(prompt, quality_select, styles_Radio, FooocusExpansion_select, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True), guidance_scale=3.5):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    generator = torch.Generator().manual_seed(seed)
    if not prompt:
        prompt = "the photo is a 18 yo jpop girl is looking absolutely adorable and gorgeous, with a playful and mischievous grin, her eyes twinkling with joy."
    if quality_select:
        prompt += ", masterpiece, best quality, very aesthetic, absurdres"
    if styles_Radio:
        for style_name in styles_Radio:
            for style in config.style_list:
                if style["name"] == style_name:
                    prompt += style["prompt"].replace("{prompt}", "the ")
    if FooocusExpansion_select:
        prompt = expansion(prompt, seed)
    image = pipe(
           prompt = "",
           prompt_2 = prompt,
           width = width,
           height = height,
           num_inference_steps = num_inference_steps, 
           generator = generator,
           guidance_scale=guidance_scale,
           output_type="pil",
    ).images[0] 
    return image, seed

def encode_image(image_path):
    """Encode the image to base64."""
    try:
        # 打开图片文件
        image = Image.open(image_path).convert("RGB")

        # 将图片转换为字节流
        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:  # 添加通用异常处理
        print(f"Error: {e}")
        return None

def predict(message, history, additional_dropdown):
    message_text = message.get("text", "")
    message_files = message.get("files", [])

    if message_files:
        # Getting the base64 string
        message_file = message_files[0]
        base64_image = encode_image(message_file)

        if base64_image is None:
            yield "Error: Failed to encode the image."
            return

        # Specify model
        model = "pixtral-large-2411"

        # Define the messages for the chat
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": message_text},
                    {
                        "type": "image_url",
                        "image_url": f"data:image/jpeg;base64,{base64_image}",
                    },
                ],
            }
        ]

        partial_message = ""
        for chunk in client.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

    else:
        stream = client_open_ai.chat.completions.create(
            model=additional_dropdown, 
            messages=[{"role": "user", "content": str(message_text)}], 
            temperature=0.5,
            max_tokens=1024,
            top_p=0.7,
            stream=True
        )

        partial_message = ""
        temp = ""
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                temp += chunk.choices[0].delta.content
                yield temp
                
with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Tab("Generator"):
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    placeholder="Enter your prompt",
                    max_lines = 12,
                    container=False
                )
                run_button = gr.Button("Run")
                result = gr.Image(label="Result", show_label=False, interactive=False)
                
                with gr.Accordion("Advanced Settings", open=False):
                
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )
                    
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    
                    with gr.Row():
                        
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=64,
                            value=896,
                        )
                        
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=64,
                            value=1152,
                        )
                    
                    with gr.Row():
                        
          
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=4,
                        )
                        guidancescale = gr.Slider(
                            label="Guidance scale",
                            minimum=0,
                            maximum=10,
                            step=0.1,
                            value=3.5,
                        )
            with gr.Tab("Styles"):
                quality_select = gr.Checkbox(label="high quality")
                FooocusExpansion_select = gr.Checkbox(label="FooocusExpansion",value=True)
                styles_name = [style["name"] for style in config.style_list]
                styles_Radio = gr.Dropdown(styles_name,label="Styles",multiselect=True)
                
        with gr.Column(scale=3,elem_id="col-container"):
            gr.ChatInterface(
                predict,
                type="messages",
                multimodal=True,
                additional_inputs =[gr.Dropdown(
                    ["CohereForAI/c4ai-command-r-plus-08-2024",
                     "meta-llama/Meta-Llama-3.1-70B-Instruct", 
                     "Qwen/Qwen2.5-72B-Instruct",
                     "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF",
                     "NousResearch/Hermes-3-Llama-3.1-8B",
                     "mistralai/Mistral-Nemo-Instruct-2407",
                     "microsoft/Phi-3.5-mini-instruct"],
                    value="meta-llama/Meta-Llama-3.1-70B-Instruct",
                    show_label=False,
                )]
            )
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [prompt, quality_select, styles_Radio, FooocusExpansion_select, seed, randomize_seed, width, height, num_inference_steps, guidancescale],
        outputs = [result, seed]
    )
if __name__ == "__main__":
    demo.queue().launch()