Update app.py

app.py

@@ -4,55 +4,79 @@ 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-72B-Instruct")
-try:
-    os.mkdir("outputs")
-except FileExistsError:
-    print("Folder already exists.")
-
-try:
-    os.mkdir("outputs/images")
-except FileExistsError:
-    print("Folder already exists.")
-
-# Define the process function that takes a text prompt and returns an image
-def process(prompt):
-    print(prompt)
+
+# Ensure output directories exist
+os.makedirs("outputs/images", exist_ok=True)
+
+# Define the process function
+def process(prompt, steps, seed, quantize, guidance, width, height):
+    print(f"Prompt: {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```\nemoji of {description}. {addon phrases}. 3D lighting. no cast shadows.\n```\n\nThe description should be a 1 sentence of your interpretation of the emoji.\nThen, 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." },
-    	{ "role": "user", "content": prompt }
+        {
+            "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."
+            ),
+        },
+        {"role": "user", "content": prompt},
     ]
 
     completion = llm.chat_completion(messages, max_tokens=100)
-
     response = completion.get("choices")[0].get("message").get("content").replace("```", "").replace("\n", "")
-    print(response)
-    time = str(datetime.now())
-    print(time)
-    time = time.replace(".", "").replace(":","")
-    image = client.text_to_image(response)
-    image.save("outputs/images/" + time + ".png")
-    with open("outputs/" + time + ".json", "w") as f:
-        f.write(json.dumps({"prompt": prompt, "refined_prompt": response, "image": "outputs/images/" + time + ".png"}))
+    print(f"Refined Prompt: {response}")
+
+    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")
+
+    with open(f"outputs/{time}.json", "w") as f:
+        json.dump({"prompt": prompt, "refined_prompt": response, "image": f"outputs/images/{time}.png"}, f)
+
     return image
 
-# Create a Gradio Blocks app
+# Create Gradio interface
 with gr.Blocks() as demo:
-    # Create a Textbox for the input prompt
-    prompt_input = gr.Textbox(label="Enter a prompt")
-    # Create an Image component for the output image
+    gr.Markdown("# Emoji Generator with Customizable Parameters")
+    
+    # Input fields
+    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)
+    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)
+    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)
+
+    # Output
     image_output = gr.Image(label="Generated Image")
-    # Create a Button to trigger the image generation
-    generate_button = gr.Button("Generate Image")
     
-    # Define the event listener for the button click
-    generate_button.click(fn=process, inputs=prompt_input, outputs=image_output)
+    # Button to generate the image
+    generate_button = gr.Button("Generate Image")
     
-    # Define the event listener for the Enter key press
-    prompt_input.submit(fn=process, inputs=prompt_input, outputs=image_output)
+    # 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,
+    )
 
-# Launch the interface
+# Launch the app
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(show_error=True)