Update app.py

app.py

@@ -125,10 +125,10 @@ if torch.cuda.is_available():
 print("Using device:", device)
 
 # --- Model Loading ---
-# Load Nanonets-OCR2-3B using AutoModelForImageTextToText
+# Load Nanonets-OCR2-3B using its specific, correct class
 MODEL_ID_V = "nanonets/Nanonets-OCR2-3B"
 processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True)
-model_v = AutoModelForImageTextToText.from_pretrained(
+model_v = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_V,
     trust_remote_code=True,
     torch_dtype=torch.float16,
@@ -179,35 +179,55 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     }]
     prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
-    # Since model is loaded with device_map="auto", we don't need to manually move inputs to device
     inputs = processor(
         text=[prompt_full],
         images=[image],
         return_tensors="pt",
-        padding=True
-    ).to(model.device)
+        padding=True).to(model.device)
 
-    # Both models now use a non-streaming generation approach
-    generation_kwargs = {
-        **inputs,
-        "max_new_tokens": max_new_tokens,
-        "do_sample": True,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "repetition_penalty": repetition_penalty,
-    }
-    
-    generated_ids = model.generate(**generation_kwargs)
-    generated_ids_trimmed = [
-        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    output_text = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-    )[0]
+    # Nanonets model supports streaming, so we use it for a better UX
+    if model_name == "Nanonets-OCR2-3B":
+        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            **inputs,
+            "streamer": streamer,
+            "max_new_tokens": max_new_tokens,
+            "do_sample": True,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repetition_penalty": repetition_penalty,
+        }
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer, buffer
     
-    output_text = output_text.replace("<|im_end|>", "").strip()
-    yield output_text, output_text
+    # Dots.OCR does not use the streamer in the same way, generate full response
+    elif model_name == "Dots.OCR":
+        generation_kwargs = {
+            **inputs,
+            "max_new_tokens": max_new_tokens,
+            "do_sample": True,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "repetition_penalty": repetition_penalty,
+        }
+        generated_ids = model.generate(**generation_kwargs)
+        generated_ids_trimmed = [
+            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+        ]
+        output_text = processor.batch_decode(
+            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        
+        output_text = output_text.replace("<|im_end|>", "").strip()
+        yield output_text, output_text
 
 
 # Define examples for image inference