feat: edit prompt using BTLM

SDXL_prompts/1.in

@@ -0,0 +1 @@
+Soft Pink Roses

SDXL_prompts/1.out

@@ -0,0 +1 @@
+Soft pink roses, white Chinese peony, tiny apple blossom flowers, eucalyptus leaves, twigs of cranberries, twigs of copper pepper berries all arrangement into a cute beautiful flowers arrangement on a pink nickel mug. The mug is sitting on a thick white book with golden cover image design. Sunny, bright image. Ad copy, huge copy space on top of the image, negative space, hd, 8k, blurry dreamy background, beautiful pictures. Product

SDXL_prompts/2.in

@@ -0,0 +1 @@
+Happy Dog

SDXL_prompts/2.out

@@ -0,0 +1 @@
+a realistic happy dog ​​playing in the grass

SDXL_prompts/3.in

@@ -0,0 +1 @@
+Gas Station

SDXL_prompts/3.out

@@ -0,0 +1,3 @@
+ours en peluche jaune , a restaurant is lit up at night on a rainy day, “gas station photography, cgsociety ), sad cop looking at a, by Nōami, picture of a loft in morning, dark ominous mood, petrol energy, a quaint, photo realistic symmetrical, drive out, pouring, terminal dark
+
+Style: Cinematic

SDXL_prompts/4.in

@@ -0,0 +1 @@
+Off Road Track

SDXL_prompts/4.out

@@ -0,0 +1,3 @@
+Prompt: A highly photorealistic image of a off road race track, complete with precise replicas of the world’s most iconic heavy noun, captured at the moment of a sharp turn, with smoke and sparks flying from under the wheels and the noun drifting around the bend. The image captures the excitement of the moment, with happy and noisy fans cheering and waving in the background. (The image is depicted at dusk, with the headlights
+
+Style: Photographic

SDXL_prompts/5.in

@@ -0,0 +1 @@
+Grapes 

SDXL_prompts/5.out

@@ -0,0 +1,3 @@
+Prompt: Powerful liquid explosion, Green grapes, Green background, Commercial photography, A bright environment, Studio lighting, OC rendering, Solid color isolated platform, Professional photography, Color classification, Super detail
+
+Style: Photographic

SDXL_prompts/6.in

@@ -0,0 +1 @@
+Astronaut

SDXL_prompts/6.out

@@ -0,0 +1 @@
+Prompt: Astronaut on Mars During sunset

SDXL_prompts/7.in

@@ -0,0 +1 @@
+Locomotive

SDXL_prompts/7.out

@@ -0,0 +1,6 @@
+Prompt: A cybernatic locomotive on rainy day from the parallel universe
+Noise: 50%
+Style realistic
+Strength 6.0
+
+Negative: Bad quality, bad detail, blurry-image, jpeg artifacts, bad contrast, bad anatomy, duplicate, watermark, extra detail, chaotic distribution of objects, distortion, bad detail facial details

main.py

@@ -0,0 +1,83 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
+
+def filter_result(result) -> str:
+    result = result.replace("\n", " ")
+    result = result[result.rfind("AGENT:") + 7:]
+    result = result.split("  ")[0]
+    result = result.split("```")[0]
+    result = result.strip()
+    return result
+
+def read_sample_prompts(path: str) -> str:
+    output = ""
+
+    # Sort files by their name to ensure they are processed in order
+    files = sorted(os.listdir(path))
+
+    # Extract and pair '.in' and '.out' files
+    input_files = [f for f in files if f.endswith('.in')]
+    output_files = [f for f in files if f.endswith('.out')]
+
+    # Ensure each input file has a corresponding output file
+    if len(input_files) != len(output_files):
+        print("Mismatch between the number of input and output files!")
+    else:
+        for idx, (in_file, out_file) in enumerate(zip(input_files, output_files)):
+            # Read input file
+            with open(os.path.join(path, in_file), 'r') as f:
+                prompt = f.read().strip()
+
+            # Read output file
+            with open(os.path.join(path, out_file), 'r') as f:
+                agent_response = f.read().strip()
+
+            # Append to the output string in the desired format
+            output += f"PROMPT: {prompt}\n"
+            output += f"AGENT: masterpiece:{idx+1}, {agent_response}\n\n"
+
+    return output
+
+
+def main(prompt_in: str) -> str:
+    tokenizer = AutoTokenizer.from_pretrained("cerebras/btlm-3b-8k-base")
+    model = AutoModelForCausalLM.from_pretrained("cerebras/btlm-3b-8k-base", trust_remote_code=True, torch_dtype="auto").to("cuda")
+
+    sample_prompts = read_sample_prompts('SDXL_prompts')
+    prompt = f"""SYSTEM: You are an agent tasked with creating better, more descriptive image prompts that result in higher quality outputs. You are given a user prompt, and determine how, if it all, to change the users image. Do you understand?
+AGENT: Yes, I understand.
+SYSTEM: Great. Below there will be a set of examples. For each example, you will be given the users prompt and the expected output. You will be then given a new user prompt and have to generate a new output. Do not output anything beyond the prompt. Do you understand?
+AGENT: Yes.
+SYSTEM: Another rule, do not list numbers in your response.
+AGENT: Yes, I understand.
+
+{sample_prompts}
+
+PROMPT: {prompt_in}
+AGENT:"""
+
+    # Change the last prompt to whatever you want to test
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+
+    # Generate text using the model
+    outputs = model.generate(
+        **inputs,
+        num_beams=5,
+        max_new_tokens=50,
+        early_stopping=True,
+        no_repeat_ngram_size=2,
+        # use random seed
+        top_k=10,
+        top_p=0.9,
+        temperature=0.9,
+        do_sample=True,
+    )
+    # Convert the generated token IDs back to text
+    generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    return generated_text[0]
+
+if __name__ == "__main__":
+    edited = main("A picture of a dog")
+    print(filter_result(edited))