File size: 5,414 Bytes
58a7875 12c9e17 58a7875 4cad52b 58a7875 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
import os
from huggingface_hub import InferenceClient
import gradio as gr
from utils import get_default_hyperparameters # Import the utility function
hf_token = os.getenv("HF_TOKEN")
headers = {
"x-wait-for-model": "true",
"x-use-cache": "true", # Use past generations
}
imagegen_models_list = [
"black-forest-labs/FLUX.1-schnell",
"black-forest-labs/FLUX.1-dev",
"strangerzonehf/Flux-Midjourney-Mix2-LoRA",
"stabilityai/stable-diffusion-3.5-large",
"stabilityai/stable-diffusion-xl-base-1.0",
"stable-diffusion-v1-5/stable-diffusion-v1-5"
]
promptgen_models_list = [
"meta-llama/Llama-3.2-11B-Vision-Instruct",
"mistralai/Mistral-Nemo-Instruct-2407"
]
# Function to generate image
def generate_image(model, prompt, guidance, width, height, num_inference_steps, seed):
print(f"""Generating image with following parameters:
Model: {model}, Prompt: {prompt}, Guidance Scale: {guidance}, Width: {width}, Height: {height},""")
client = InferenceClient(model=model, headers=headers, token=hf_token)
image = client.text_to_image(
model=model,
prompt=prompt,
guidance_scale=guidance,
height=height,
width=width,
num_inference_steps=num_inference_steps,
seed=seed
)
return image
# Function to update hyperparameters dynamically
def update_hyperparameters(model_name):
default_params = get_default_hyperparameters(model_name)
return (default_params['guidance_scale'],
default_params['width'],
default_params['height'],
default_params['num_inference_steps'],
default_params['seed'])
# Function to expand the idea using the selected prompt generation model
def expand_idea(promptgen_model, idea_text):
print(f"Expanding idea with model: {promptgen_model}")
client = InferenceClient(model=promptgen_model, headers=headers, token=hf_token)
response = client.chat_completion(
messages=[
{
'role': 'user',
'content': f'For the given idea, generate a text prompt to generate an image from a text to image generator. Be creative and include both subject and style prompts into one. Do not, explain your decisions. Idea: {idea_text}',
},
],
max_tokens=100,
temperature=1.1, # Set temperature higher for dynamic responses
top_p=0.9,
)
expanded_prompt = response.choices[0].message.content # Assuming this is how the prompt is expanded
return expanded_prompt
# Interface for generating images and expanding idea into a prompt
def run_interface():
with gr.Blocks() as iface:
with gr.Row():
# Image generation controls
with gr.Column(scale=1):
model_dropdown = gr.Dropdown(choices=imagegen_models_list, label="Image Model", value=imagegen_models_list[0])
prompt_textbox = gr.Textbox(label="Prompt", lines=5, value="Astronaut floating in space")
# Initial default values based on the first model
default_params = get_default_hyperparameters(model_dropdown.value)
guidance_slider = gr.Slider(0, 10, step=0.1, label="Guidance Scale", value=default_params['guidance_scale'])
width_slider = gr.Slider(256, 2048, step=32, label="Width", value=default_params['width'])
height_slider = gr.Slider(256, 2048, step=32, label="Height", value=default_params['height'])
steps_slider = gr.Slider(1, 100, step=1, label="Number of Inference Steps", value=default_params['num_inference_steps'])
seed_number = gr.Number(label="Seed", value=default_params['seed'])
# Update sliders based on model selection
model_dropdown.change(
fn=update_hyperparameters,
inputs=model_dropdown,
outputs=[guidance_slider, width_slider, height_slider, steps_slider, seed_number]
)
generate_button = gr.Button("Generate Image")
output_image = gr.Image(type="pil", format="png")
generate_button.click(
fn=generate_image,
inputs=[model_dropdown, prompt_textbox, guidance_slider, width_slider, height_slider, steps_slider, seed_number],
outputs=output_image
)
# Prompt expansion section in a collapsible panel
with gr.Column(scale=1):
gr.Markdown("**Expand an Idea into Prompt**")
with gr.Accordion("Your Image Idea", open=False):
promptgen_dropdown = gr.Dropdown(choices=promptgen_models_list, label="Prompt Generation Model", value=promptgen_models_list[0])
idea_textbox = gr.Textbox(label="Enter your idea/sketch (rough text)", lines=5, placeholder="Describe your idea here...")
expand_button = gr.Button("Expand to Prompt")
# Expand idea and directly output to the main prompt textbox
expand_button.click(
fn=expand_idea,
inputs=[promptgen_dropdown, idea_textbox],
outputs=prompt_textbox # Updating the prompt_textbox directly
)
iface.launch()
run_interface() |