Update raw.py

raw.py

@@ -80,32 +80,22 @@ pipe.to("cuda")
 @spaces.GPU()
 @torch.no_grad()
 def caption(input_image: Image.Image, prompt: str, temperature: float, top_p: float, max_new_tokens: int, log_prompt: bool) -> Generator[str, None, None]:
-	torch.cuda.empty_cache()
-
-	if input_image is None:
-		yield "No image provided. Please upload an image."
-		return
-
-	if log_prompt:
-		print(f"PromptLog: {repr(prompt)}")
-
-	convo = [
-		{
-			"role": "system",
-			# Beta One supports a wider range of system prompts, but this is a good default
-			"content": "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions.",
-		},
-		{
-			"role": "user",
-			"content": prompt.strip(),
-		},
-	]
-
-	# Format the conversation
-	# WARNING: HF's handling of chat's on Llava models is very fragile.  This specific combination of processor.apply_chat_template(), and processor() works
-	# but if using other combinations always inspect the final input_ids to ensure they are correct.  Often times you will end up with multiple <bos> tokens
-	# if not careful, which can make the model perform poorly.
-    
+    torch.cuda.empty_cache()
+    if input_image is None:
+        yield "No image provided. Please upload an image."
+        return
+    if log_prompt:
+        print(f"PromptLog: {repr(prompt)}")
+    convo = [
+        {
+            "role": "system",
+            "content": "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions.",
+        },
+        {
+            "role": "user",
+            "content": prompt.strip(),
+        },
+    ]
     convo_string = cap_processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
     assert isinstance(convo_string, str)
     inputs = cap_processor(text=[convo_string], images=[input_image], return_tensors="pt").to('cuda')
@@ -127,6 +117,7 @@ def caption(input_image: Image.Image, prompt: str, temperature: float, top_p: fl
     for text in streamer:
         outputs.append(text)
         yield "".join(outputs)
+
         
 @spaces.GPU()
 @torch.no_grad()