Update app.py

app.py

@@ -1,245 +1,121 @@
-import gradio as gr
 import torch
-import requests
-from transformers import AutoModel, AutoTokenizer
-import spaces
-from typing import Iterable
+from transformers import AutoTokenizer, AutoModel
+from huggingface_hub import HfApi, HfFolder, create_repo, upload_folder
 import os
-import tempfile
-from PIL import Image, ImageDraw
-import re
-from gradio.themes import Soft
-from gradio.themes.utils import colors, fonts, sizes
-from docling_core.types.doc import DoclingDocument, DocTagsDocument
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
-print("torch.__version__ =", torch.__version__)
-print("torch.version.cuda =", torch.version.cuda)
-print("cuda available:", torch.cuda.is_available())
-print("cuda device count:", torch.cuda.device_count())
-if torch.cuda.is_available():
-    print("current device:", torch.cuda.current_device())
-    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
-
-print("Using device:", device)
-
-
-colors.steel_blue = colors.Color(
-    name="steel_blue",
-    c50="#EBF3F8",
-    c100="#D3E5F0",
-    c200="#A8CCE1",
-    c300="#7DB3D2",
-    c400="#529AC3",
-    c500="#4682B4",
-    c600="#3E72A0",
-    c700="#36638C",
-    c800="#2E5378",
-    c900="#264364",
-    c950="#1E3450",
-)
-
-class SteelBlueTheme(Soft):
-    def __init__(
-        self,
-        *,
-        primary_hue: colors.Color | str = colors.gray,
-        secondary_hue: colors.Color | str = colors.steel_blue,
-        neutral_hue: colors.Color | str = colors.slate,
-        text_size: sizes.Size | str = sizes.text_lg,
-        font: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
-        ),
-        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
-            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
-        ),
-    ):
-        super().__init__(
-            primary_hue=primary_hue,
-            secondary_hue=secondary_hue,
-            neutral_hue=neutral_hue,
-            text_size=text_size,
-            font=font,
-            font_mono=font_mono,
-        )
-        super().set(
-            background_fill_primary="*primary_50",
-            background_fill_primary_dark="*primary_900",
-            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
-            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
-            button_primary_text_color="white",
-            button_primary_text_color_hover="white",
-            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
-            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
-            slider_color="*secondary_500",
-            slider_color_dark="*secondary_600",
-            block_title_text_weight="600",
-            block_border_width="3px",
-            block_shadow="*shadow_drop_lg",
-            button_primary_shadow="*shadow_drop_lg",
-            button_large_padding="11px",
-            color_accent_soft="*primary_100",
-            block_label_background_fill="*primary_200",
-        )
+import gradio as gr
 
-steel_blue_theme = SteelBlueTheme()
-
-css = """
-#main-title h1 {
-    font-size: 2.3em !important;
-}
-#output-title h2 {
-    font-size: 2.1em !important;
-}
-"""
-
-print("Determining device...")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"✅ Using device: {device}")
-
-print("Loading model and tokenizer...")
-model_name = "prithivMLmods/DeepSeek-OCR-transformers-5.0.0.dev0" # -> Latest transformers version used for the model. (https://huggingface.co/deepseek-ai/DeepSeek-OCR)
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-
-model = AutoModel.from_pretrained(
-    model_name,
-    #_attn_implementation="flash_attention_2",
-    trust_remote_code=True,
-    use_safetensors=True,
-).to(device).eval() # Move to device and set to eval mode
-
-if device.type == 'cuda':
-    model = model.to(torch.bfloat16)
-
-print("✅ Model loaded successfully to device and in eval mode.")
-
-def find_result_image(path):
-    for filename in os.listdir(path):
-        if "grounding" in filename or "result" in filename:
-            try:
-                image_path = os.path.join(path, filename)
-                return Image.open(image_path)
-            except Exception as e:
-                print(f"Error opening result image {filename}: {e}")
-    return None
-
-@spaces.GPU
-def process_ocr_task(image, model_size, task_type, ref_text):
+def load_and_reupload_model(model_name, new_repo_id, hf_token, max_shard_size="1.5GB"):
     """
-    Processes an image with DeepSeek-OCR. The model is already on the correct device.
+    Loads a model and tokenizer, saves the model in smaller shards,
+    and uploads them to a new private repository on the Hugging Face Hub.
+
+    Args:
+        model_name (str): The name of the model to load from the Hugging Face Hub.
+        new_repo_id (str): The ID for the new private repository (e.g., "your-username/your-repo-name").
+        hf_token (str): Your Hugging Face API token with write permissions.
+        max_shard_size (str, optional): The maximum size of each model shard. Defaults to "1GB".
+    
+    Returns:
+        A string log of the process.
     """
-    if image is None:
-        return "Please upload an image first.", None
-
-    print("✅ Model is already on the designated device.")
-
-    with tempfile.TemporaryDirectory() as output_path:
-        # Build the prompt
-        if task_type == "Free OCR":
-            prompt = "<image>\nFree OCR."
-        elif task_type == "Convert to Markdown":
-            prompt = "<image>\n<|grounding|>Convert the document to markdown."
-        elif task_type == "Parse Figure":
-            prompt = "<image>\nParse the figure."
-        elif task_type == "Locate Object by Reference":
-            if not ref_text or ref_text.strip() == "":
-                raise gr.Error("For the 'Locate' task, you must provide the reference text to find!")
-            prompt = f"<image>\nLocate <|ref|>{ref_text.strip()}<|/ref|> in the image."
-        else:
-            prompt = "<image>\nFree OCR."
-
-        temp_image_path = os.path.join(output_path, "temp_image.png")
-        image.save(temp_image_path)
-
-        size_configs = {
-            "Tiny": {"base_size": 512, "image_size": 512, "crop_mode": False},
-            "Small": {"base_size": 640, "image_size": 640, "crop_mode": False},
-            "Base": {"base_size": 1024, "image_size": 1024, "crop_mode": False},
-            "Large": {"base_size": 1280, "image_size": 1280, "crop_mode": False},
-            "Gundam (Recommended)": {"base_size": 1024, "image_size": 640, "crop_mode": True},
-        }
-        config = size_configs.get(model_size, size_configs["Gundam (Recommended)"])
-
-        print(f"🏃 Running inference with prompt: {prompt}")
-        text_result = model.infer(
-            tokenizer,
-            prompt=prompt,
-            image_file=temp_image_path,
-            output_path=output_path,
-            base_size=config["base_size"],
-            image_size=config["image_size"],
-            crop_mode=config["crop_mode"],
-            save_results=True,
-            test_compress=True,
-            eval_mode=True,
+    log_output = []
+    try:
+        # --- Validate inputs ---
+        if not all([model_name, new_repo_id, hf_token]):
+            return "❌ Error: All fields (Model Name, New Repo ID, and HF Token) are required."
+
+        # --- 1. Log in to Hugging Face Hub ---
+        HfFolder.save_token(hf_token)
+        log_output.append("✅ Successfully saved Hugging Face token.")
+
+        # --- 2. Create a new private repository ---
+        create_repo(new_repo_id, private=True, exist_ok=True, token=hf_token)
+        log_output.append(f"✅ Successfully created or confirmed private repository: {new_repo_id}")
+
+        # --- 3. Load the tokenizer and model ---
+        log_output.append(f"🔄 Loading tokenizer for {model_name}...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        log_output.append("✅ Tokenizer loaded successfully.")
+
+        log_output.append(f"🔄 Loading model {model_name}...")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = AutoModel.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            use_safetensors=True,
+        ).to(device)
+
+        if device.type == 'cuda':
+            model = model.to(torch.bfloat16)
+
+        model.eval()
+        log_output.append(f"✅ Model loaded successfully to {device} and set to evaluation mode.")
+
+        # --- 4. Save the model and tokenizer locally with sharding ---
+        local_save_dir = new_repo_id.split('/')[-1]
+        os.makedirs(local_save_dir, exist_ok=True)
+
+        log_output.append(f"🔄 Saving model with max_shard_size='{max_shard_size}' to '{local_save_dir}'...")
+        model.save_pretrained(local_save_dir, max_shard_size=max_shard_size)
+        tokenizer.save_pretrained(local_save_dir)
+        log_output.append("✅ Model and tokenizer saved locally.")
+
+        # --- 5. Upload the sharded model and tokenizer to the new repo ---
+        log_output.append(f"🔄 Uploading files to {new_repo_id}...")
+        api = HfApi()
+        api.upload_folder(
+            folder_path=local_save_dir,
+            repo_id=new_repo_id,
+            repo_type="model",
+            token=hf_token
         )
+        log_output.append(f"🚀 Successfully uploaded model and tokenizer to private repo: {new_repo_id}")
 
-        print(f"====\n📄 Text Result: {text_result}\n====")
-
-        result_image_pil = None
-        pattern = re.compile(r"<\|det\|>\[\[(\d+),\s*(\d+),\s*(\d+),\s*(\d+)\]\]<\|/det\|>")
-        matches = list(pattern.finditer(text_result))
-
-        if matches:
-            print(f"✅ Found {len(matches)} bounding box(es). Drawing on the original image.")
-            image_with_bboxes = image.copy()
-            draw = ImageDraw.Draw(image_with_bboxes)
-            w, h = image.size
-
-            for match in matches:
-                coords_norm = [int(c) for c in match.groups()]
-                x1_norm, y1_norm, x2_norm, y2_norm = coords_norm
-
-                x1 = int(x1_norm / 1000 * w)
-                y1 = int(y1_norm / 1000 * h)
-                x2 = int(x2_norm / 1000 * w)
-                y2 = int(y2_norm / 1000 * h)
-
-                draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
-
-            result_image_pil = image_with_bboxes
-        else:
-            print("⚠️ No bounding box coordinates found in text result. Falling back to search for a result image file.")
-            result_image_pil = find_result_image(output_path)
-
-        return text_result, result_image_pil
-
-# url = "https://huggingface.co/spaces/prithivMLmods/Multimodal-OCR3/resolve/main/examples/3.jpg?download=true"
-# example_image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+    except Exception as e:
+        log_output.append(f"❌ An error occurred: {e}")
+    
+    return "\n".join(log_output)
 
-with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
-    gr.Markdown("# **DeepSeek OCR [exp]**", elem_id="main-title")
+# --- Gradio Interface ---
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown(
+        """
+        # 🚀 Hugging Face Model Sharder & Re-Uploader
+        This application loads a model from the Hugging Face Hub, saves it locally into smaller shards (e.g., 1GB each), and then uploads it to a new private repository under your account.
+        """
+    )
 
-    
     with gr.Row():
-        with gr.Column(scale=1):
-            image_input = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard"])
-            model_size = gr.Dropdown(choices=["Tiny", "Small", "Base", "Large", "Gundam (Recommended)"], value="Large", label="Resolution Size")
-            task_type = gr.Dropdown(choices=["Free OCR", "Convert to Markdown", "Parse Figure", "Locate Object by Reference"], value="Convert to Markdown", label="Task Type")
-            ref_text_input = gr.Textbox(label="Reference Text (for Locate task)", placeholder="e.g., the teacher, 20-10, a red car...", visible=False)
-            submit_btn = gr.Button("Process Image", variant="primary")
-
-            examples = gr.Examples(
-                examples=["examples/1.jpg", "examples/2.jpg", "examples/3.jpg"],
-                                inputs=image_input, label="Examples"
-                            )
-
         with gr.Column(scale=2):
-            output_text = gr.Textbox(label="Output (OCR)", lines=8, show_copy_button=True)
-            output_image = gr.Image(label="Layout Detection (If Any)", type="pil")
-            
-            with gr.Accordion("Note", open=False):
-                gr.Markdown("Inference using Huggingface transformers on NVIDIA GPUs. This app is running with transformers version 4.57.1 and torch version 2.6.0.")
-                
-    def toggle_ref_text_visibility(task):
-        return gr.Textbox(visible=True) if task == "Locate Object by Reference" else gr.Textbox(visible=False)
-
-    task_type.change(fn=toggle_ref_text_visibility, inputs=task_type, outputs=ref_text_input)
-    submit_btn.click(fn=process_ocr_task, inputs=[image_input, model_size, task_type, ref_text_input], outputs=[output_text, output_image])
+            model_name_input = gr.Textbox(
+                label="Original Model Name",
+                value="strangervisionhf/deepseek-ocr-latest-transformers",
+                placeholder="e.g., 'strangervisionhf/deepseek-ocr-latest-transformers'"
+            )
+            new_repo_id_input = gr.Textbox(
+                label="New Private Repository ID",
+                placeholder="e.g., 'your-username/private-deepseek-ocr-sharded'"
+            )
+            hf_token_input = gr.Textbox(
+                label="Hugging Face Write Token",
+                type="password",
+                placeholder="Enter your Hugging Face token with write access"
+            )
+            run_button = gr.Button("Shard and Upload Model", variant="primary")
+        
+        with gr.Column(scale=3):
+            output_log = gr.Textbox(
+                label="Process Log",
+                lines=15,
+                interactive=False,
+                autoscroll=True
+            )
+
+    run_button.click(
+        fn=load_and_reupload_model,
+        inputs=[model_name_input, new_repo_id_input, hf_token_input],
+        outputs=output_log
+    )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(share=True, mcp_server=True, ssr_mode=False)
+    demo.launch()