Initial Commits

text_image.py

@@ -0,0 +1,124 @@
+import os
+from dotenv import load_dotenv, find_dotenv
+from huggingface_hub import InferenceClient
+import gradio as gr
+
+from utils import get_default_hyperparameters  # Import the utility function
+
+load_dotenv(find_dotenv())
+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=80,
+        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()

utils.py

@@ -0,0 +1,66 @@
+# util.py
+
+def get_default_hyperparameters(model_name):
+    """
+    Returns default hyperparameters based on the model name.
+    
+    Args:
+        model_name (str): Name of the selected model.
+
+    Returns:
+        dict: A dictionary of default hyperparameters.
+    """
+    # Define default hyperparameters for each model
+    default_hyperparameters = {
+        "black-forest-labs/FLUX.1-schnell": {
+            "guidance_scale": 0,
+            "width": 1024,
+            "height": 1024,
+            "num_inference_steps": 4,
+            "seed": 0
+        },
+        "black-forest-labs/FLUX.1-dev": {
+            "guidance_scale": 3.5,
+            "width": 1024,
+            "height": 1024,
+            "num_inference_steps": 28,
+            "seed": 0
+        },
+        "strangerzonehf/Flux-Midjourney-Mix2-LoRA": {
+            "guidance_scale": 3.5,
+            "width": 1024,
+            "height": 1024,
+            "num_inference_steps": 28,
+            "seed": 0
+        },
+        "stabilityai/stable-diffusion-3.5-large": {
+            "guidance_scale": 4.5,
+            "width": 1024,
+            "height": 1024,
+            "num_inference_steps": 35,
+            "seed": 0
+        },
+        "stabilityai/stable-diffusion-xl-base-1.0": {
+            "guidance_scale": 7,
+            "width": 1024,
+            "height": 1024,
+            "num_inference_steps": 30,
+            "seed": 0
+        },
+        "stable-diffusion-v1-5/stable-diffusion-v1-5": {
+            "guidance_scale": 5.0,
+            "width": 512,
+            "height": 512,
+            "num_inference_steps": 20,
+            "seed": 0
+        }
+    }
+    
+    # Return the hyperparameters for the selected model or a default set if not found
+    return default_hyperparameters.get(model_name, {
+        "guidance_scale": 7.5,
+        "width": 512,
+        "height": 512,
+        "num_inference_steps": 20,
+        "seed": 42
+    })