Reality123b/iris

1import gradio as gr  2from huggingface_hub import InferenceClient  3from PIL import Image  4import time  5import os  6import base64  7from io import BytesIO  8 9HF_TOKEN = os.environ.get("HF_TOKEN")  10 11if not HF_TOKEN:  12    HF_TOKEN_ERROR = "Hugging Face API token (HF_TOKEN) not found. Please set it as an environment variable or Gradio secret."  13else:  14    HF_TOKEN_ERROR = None  15 16client = InferenceClient(token=HF_TOKEN)  17PROMPT_IMPROVER_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"  18 19def improve_prompt(original_prompt):  20    if HF_TOKEN_ERROR:  21        raise gr.Error(HF_TOKEN_ERROR)  22 23    try:  24        system_prompt = "You are a helpful assistant that improves text prompts for image generation models. Make the prompt more descriptive, detailed, and artistic, while keeping the user's original intent."  25        prompt_for_llm = f"""<|system|>  26{system_prompt}</s>  27<|user|>  28Improve this prompt: {original_prompt}  29</s>  30<|assistant|>  31"""  32        improved_prompt = client.text_generation(  33            prompt=prompt_for_llm,  34            model=PROMPT_IMPROVER_MODEL,  35            max_new_tokens=1280,  36            temperature=0.7,  37            top_p=0.9,  38            repetition_penalty=1.2,  39            stop_sequences=["</s>"],  40        )  41 42        return improved_prompt.strip()  43 44    except Exception as e:  45        print(f"Error improving prompt: {e}")  46        return original_prompt  47 48 49def generate_image(prompt, progress=gr.Progress()):  50    if HF_TOKEN_ERROR:  51        raise gr.Error(HF_TOKEN_ERROR)  52 53    progress(0, desc="Improving prompt...")  54    improved_prompt = improve_prompt(prompt)  55 56    progress(0.2, desc="Sending request ")  57    try:  58        image = client.text_to_image(improved_prompt, model="black-forest-labs/FLUX.1-schnell")  59 60        if not isinstance(image, Image.Image):  61            raise Exception(f"Expected a PIL Image, but got: {type(image)}")  62 63        progress(0.8, desc="Processing image...")  64        time.sleep(0.5)  65        progress(1.0, desc="Done!")  66        return image  67    except Exception as e:  68        if "rate limit" in str(e).lower():  69            error_message = f"Rate limit exceeded. Please try again later. Error: {e}"  70        else:  71            error_message = f"An error occurred: {e}"  72        raise gr.Error(error_message)  73 74 75def pil_to_base64(img):  76    buffered = BytesIO()  77    img.save(buffered, format="PNG")  78    img_str = base64.b64encode(buffered.getvalue()).decode()  79    return f"data:image/png;base64,{img_str}"  80 81 82css = """  83 84"""  85 86with gr.Blocks(css=css) as demo:  87    gr.Markdown(  88        """  89        # Xylaria Iris v3  90        """,  91        elem_classes="title"  92    )  93 94    with gr.Row():  95        with gr.Column():  96            with gr.Group(elem_classes="input-section"):  97                prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., A cat", lines=3)  98                generate_button = gr.Button("Generate Image", elem_classes="submit-button")  99        with gr.Column():  100            with gr.Group(elem_classes="output-section") as output_group:  101                image_output = gr.Image(label="Generated Image", interactive=False)  102 103    def on_generate_click(prompt):  104        output_group.elem_classes = ["output-section", "animate"]  105        image = generate_image(prompt)  # Ignore the improved prompt  106        output_group.elem_classes = ["output-section"]  107 108        return image  # Return only the generated image  109 110    generate_button.click(on_generate_click, inputs=prompt_input, outputs=image_output)  111    prompt_input.submit(on_generate_click, inputs=prompt_input, outputs=image_output)  112 113    gr.Examples(  114        [["A dog"],  115         ["A house on a hill"],  116         ["A spaceship"]],  117        inputs=prompt_input  118    )  119 120if __name__ == "__main__":  121    demo.queue().launch()