Update app.py

app.py

@@ -119,28 +119,28 @@ def king(type = "Image Editing",
         return seed, output_image
         
 def response(instruction, input_image=None):
-    client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-        
-    generate_kwargs = dict(
-            max_new_tokens=5,
-        )
-
-    system="[SYSTEM] You will be provided with text, and your task is to classify task is image generation or image editing answer with only task do not say anything else and stop as soon as possible. [TEXT]"
-        
-    formatted_prompt = system + instruction + "[TASK]"
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-    for response in stream:        
-        if not response.token.text == "</s>":
-            output += response.token.text
-
     if input_image is None:
         output="Image Generation"
-        
-    if "editing" in output:
-        output = "Image Editing"
     else:
-        output = "Image Generation"
+        client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+            
+        generate_kwargs = dict(
+                max_new_tokens=5,
+            )
+    
+        system="[SYSTEM] You will be provided with text, and your task is to classify task is image generation or image editing answer with only task do not say anything else and stop as soon as possible. [TEXT]"
+            
+        formatted_prompt = system + instruction + "[TASK]"
+        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+        output = ""
+        for response in stream:        
+            if not response.token.text == "</s>":
+                output += response.token.text
+        if "editing" in output:
+            output = "Image Editing"
+        else:
+            output = "Image Generation"
+        yield output
     return output
 
 css = '''
@@ -152,31 +152,36 @@ footer {
 '''
 
 def get_example():
-    case = [
+ case = [
         [
             "Image Generation",
-            "A Super Car",
             None,
+            "A Super Car",
+
         ],
         [
             "Image Editing",
-            "make it red",
             "./supercar.png",
+            "make it red",
+
         ],
         [
             "Image Editing",
-            "add some snow",
             "./red_car.png",
+            "add some snow",
+
         ],
         [
             "Image Generation",
-            "Ironman flying in front of Ststue of liberty",
             None,
+            "Kids going o school, Anime style",
+
         ],
         [
             "Image Generation",
+            None,
             "Beautiful Eiffel Tower at Night",
-             None,
+
         ],
     ]
     return case