test mistral + flux

app.py

@@ -4,9 +4,7 @@ import random
 import spaces
 import torch
 from diffusers import DiffusionPipeline
-from huggingface_hub import InferenceClient
-
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+from transformers import pipeline
 
 
 dtype = torch.bfloat16
@@ -17,6 +15,19 @@ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", tor
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+
+
+
+def refine_prompt(prompt):
+    chatbot = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.3")
+    messages = [
+        {"role": "system", "content": "You are a product designer. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Produce an extended description of product front view that will be use by Flux to generate a visual"},
+        {"role": "user", "content": "a castle schoolbag"},
+    ]
+    chatbot(messages)
+    refined_prompt = chatbot(prompt)
+    return refined_prompt
+
 @spaces.GPU()
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
@@ -62,7 +73,7 @@ with gr.Blocks(css=css) as demo:
                 container=False,
             )
             
-            run_button = gr.Button("Run", scale=0)
+            prompt_button = gr.Button("Refine prompt", scale=0)
         
         refined_prompt = gr.Text(
             label="Refined Prompt", 
@@ -71,10 +82,13 @@ with gr.Blocks(css=css) as demo:
             placeholder="Prompt refined by Mistral", 
             container=False
             )
+        
+        
+        run_button = gr.Button("Create visual", scale=0)
 
         result = gr.Image(label="Result", show_label=False)
         
-        with gr.Accordion("Advanced Settings istral", open=False):
+        with gr.Accordion("Advanced Settings Mistral", open=False):
             gr.Slider(
                 label="Temperature",
                 value=0.9,
@@ -161,10 +175,18 @@ with gr.Blocks(css=css) as demo:
             cache_examples="lazy"
         )
 
+
+    gr.on(
+        triggers=[prompt_button.click, prompt.submit],
+        fn = refine_prompt,
+        inputs = [prompt],
+        outputs = [refined_prompt]
+    )
+
     gr.on(
-        triggers=[run_button.click, prompt.submit],
+        triggers=[run_button.click],
         fn = infer,
-        inputs = [prompt, seed, randomize_seed, width, height, num_inference_steps],
+        inputs = [refined_prompt, seed, randomize_seed, width, height, num_inference_steps],
         outputs = [result, seed]
     )