Update app.py

app.py

@@ -1,80 +1,107 @@
 import gradio as gr
+import torch
+from diffusers import FluxPipeline
 from huggingface_hub import InferenceClient
 import os
-from datetime import datetime
-import json
 
-# Initialize clients
-client = InferenceClient("EvanZhouDev/open-genmoji", token=os.getenv("HUGGINGFACE_API_TOKEN"))
-llm = InferenceClient("Qwen/Qwen2.5-3B-Instruct")
+# Initialize the Flux pipeline
+def initialize_flux_pipeline():
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
+        torch_dtype=torch.bfloat16
+    )
+    pipe.enable_model_cpu_offload()
+    pipe.load_lora_weights("EvanZhouDev/open-genmoji")
+    return pipe
 
-# Ensure output directories exist
-os.makedirs("outputs/images", exist_ok=True)
+flux_pipeline = initialize_flux_pipeline()
 
-# Define the process function
-def process(prompt, steps, seed, quantize, guidance, width, height):
-    print(f"Prompt: {prompt}")
+# Initialize the language model client
+llm_client = InferenceClient("Qwen/Qwen2.5-72B-Instruct", token=os.getenv("HUGGINGFACE_API_TOKEN"))
+
+# Function to refine the prompt
+def refine_prompt(original_prompt):
     messages = [
         {
             "role": "system",
             "content": (
-                "You are helping create a prompt for a Emoji generation image model. An emoji must be easily "
-                "interpreted when small so details must be exaggerated to be clear. Your goal is to use descriptions "
-                "to achieve this.\n\nYou will receive a user description, and you must rephrase it to consist of "
-                "short phrases separated by periods, adding detail to everything the user provides.\n\nAdd describe "
-                "the color of all parts or components of the emoji. Unless otherwise specified by the user, do not "
-                "describe people. Do not describe the background of the image. Your output should be in the format:\n\n"
-                "```emoji of {description}. {addon phrases}. 3D lighting. no cast shadows.```\n\nThe description "
-                "should be a 1 sentence of your interpretation of the emoji. Then, you may choose to add addon phrases."
-                " You must use the following in the given scenarios:\n\n- \"cute.\": If generating anything that's not "
-                "an object, and also not a human\n- \"enlarged head in cartoon style.\": ONLY animals\n- \"head is "
-                "turned towards viewer.\": ONLY humans or animals\n- \"detailed texture.\": ONLY objects\n\nFurther "
-                "addon phrases may be added to ensure the clarity of the emoji."
+                "You are refining a user-provided description for generating images. The output should focus on clarity, "
+                "detail, and vivid descriptions. The format should be concise and effective for image generation."
             ),
         },
-        {"role": "user", "content": prompt},
+        {"role": "user", "content": original_prompt},
     ]
+    completion = llm_client.chat_completion(messages, max_tokens=100)
+    refined = completion["choices"][0]["message"]["content"].strip()
+    return refined
 
-    completion = llm.chat_completion(messages, max_tokens=100)
-    response = completion.get("choices")[0].get("message").get("content").replace("```", "").replace("\n", "")
-    print(f"Refined Prompt: {response}")
+# Define the process function
+def process(prompt, guidance_scale, num_inference_steps, height, width, seed):
+    print(f"Original Prompt: {prompt}")
 
-    time = datetime.now().strftime("%Y%m%d%H%M%S")
-    image = client.text_to_image(response, steps=steps, seed=seed, quantize=quantize, guidance=guidance, width=width, height=height)
-    image.save(f"outputs/images/{time}.png")
+    # Refine the prompt
+    try:
+        refined_prompt = refine_prompt(prompt)
+        print(f"Refined Prompt: {refined_prompt}")
+    except Exception as e:
+        return f"Error refining prompt: {str(e)}"
 
-    with open(f"outputs/{time}.json", "w") as f:
-        json.dump({"prompt": prompt, "refined_prompt": response, "image": f"outputs/images/{time}.png"}, f)
+    # Set the random generator seed
+    generator = torch.Generator(device="cuda").manual_seed(seed)
 
-    return image
+    try:
+        # Generate the image
+        output = flux_pipeline(
+            prompt=refined_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            height=height,
+            width=width,
+            generator=generator,
+        )
+        image = output.images[0]
+        return image
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
 
-# Create Gradio interface
+# Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Emoji Generator with Customizable Parameters")
-    
-    # Input fields
+    gr.Markdown("# Flux Text-to-Image Generator with Prompt Refinement")
+
+    # User inputs
     with gr.Row():
-        prompt_input = gr.Textbox(label="Enter a prompt")
-        steps_input = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1)
+        prompt_input = gr.Textbox(label="Enter a Prompt", placeholder="Describe your image")
+        guidance_scale_input = gr.Slider(
+            label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.5, step=0.1
+        )
     with gr.Row():
-        seed_input = gr.Number(label="Seed", value=2, precision=0)
-        quantize_input = gr.Slider(label="Quantize", minimum=1, maximum=16, value=8, step=1)
+        num_inference_steps_input = gr.Slider(
+            label="Inference Steps", minimum=1, maximum=100, value=50, step=1
+        )
+        seed_input = gr.Number(label="Seed", value=42, precision=0)
     with gr.Row():
-        guidance_input = gr.Slider(label="Guidance", minimum=1.0, maximum=10.0, value=5.0, step=0.1)
-        width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=1280, step=64)
-        height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=640, step=64)
+        height_input = gr.Slider(label="Height", minimum=256, maximum=2048, value=768, step=64)
+        width_input = gr.Slider(label="Width", minimum=256, maximum=2048, value=1360, step=64)
 
-    # Output
+    # Output components
+    refined_prompt_output = gr.Textbox(label="Refined Prompt", interactive=False)
     image_output = gr.Image(label="Generated Image")
-    
+
     # Button to generate the image
     generate_button = gr.Button("Generate Image")
-    
+
     # Define button click behavior
     generate_button.click(
-        fn=process,
-        inputs=[prompt_input, steps_input, seed_input, quantize_input, guidance_input, width_input, height_input],
-        outputs=image_output,
+        fn=lambda prompt, *args: (refine_prompt(prompt), process(prompt, *args)),
+        inputs=[
+            prompt_input,
+            guidance_scale_input,
+            num_inference_steps_input,
+            height_input,
+            width_input,
+            seed_input,
+        ],
+        outputs=[refined_prompt_output, image_output],
     )
 
 # Launch the app