Update app.py

app.py

@@ -15,7 +15,7 @@ repo = "stabilityai/stable-diffusion-3-medium"
 t2i = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16, revision="refs/pr/26",token=os.environ["TOKEN"]).to(device)
 
 model_id = "microsoft/Phi-3-medium-128k-instruct"
-model = AutoModelForCausalLM.from_pretrained(
+upsampler = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map=device, 
     torch_dtype=torch.bfloat16, 
@@ -23,19 +23,6 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-upsampler = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-)
-
-generation_args = {
-    "max_new_tokens": 300,
-    "return_full_text": False,
-    "temperature": 0.7,
-    "do_sample": True,
-}
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1344
 
@@ -54,8 +41,18 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         {"role": "assistant", "content": "A gourmet scene in a high-end restaurant kitchen where a chef is presenting a plate of cooked beef testicles, garnished elegantly with herbs and spices. The chef, a middle-aged Caucasian man wearing a white chef's hat and coat, is inspecting the dish with a satisfied expression. The kitchen background is bustling with other chefs and kitchen staff, and the atmosphere is warm and inviting with hanging pots and pans, and a glowing, busy stove in the background. The focus is on the chef's proud presentation of this unusual but delicately prepared dish."},
         {"role": "user", "content": prompt},
     ]
-    output = upsampler(messages, **generation_args)
-    upsampled_prompt=output[0]['generated_text']
+    tokenized_input = tokenizer.apply_chat_templete(messages, add_special_tokens=False, add_generate_prompt=True, return_tensors="pt").to(model.device)
+    with torch.inference_mode():
+        output = upsampler.generate(
+            tokenized_input,
+            max_new_tokens=512,
+            do_sample=True,
+            top_p=0.95,
+            temperature=0.7,
+            repetition_penalty=1.05,
+        )[0]
+    print(tokenizer.decode(output))
+    upsampled_prompt=output['generated_text']
 
     print(upsampled_prompt)