Update app.py

app.py

@@ -96,30 +96,19 @@ def king(type ,
         return seed, image
 
 # Prompt classifier
-def response(instruction, input_image=None):
+def response(instruction, input_image=None ):
     if input_image is None:
         output="Image Generation"
-        yield output
     else:
-        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:
+        text = instruction
+        labels = ["Image Editing", "Image Generation"]
+        classification = client.zero_shot_classification(text, labels, multi_label=True)
+        output = classification[0]
+        output = str(output)
+        if "Editing" in output:
             output = "Image Editing"
         else:
             output = "Image Generation"
-        yield output
     return output
 
 css = '''