Ruian7P/imuru

1import gradio as gr2import spaces3import torch4import json5from pathlib import Path6from PIL import Image7import numpy as np8 9def ensure_env_installed():10    try:11        import transformers12        import torchvision13        import diffusers14        import einops15    except ImportError:16        import subprocess17        import sys18        subprocess.check_call([sys.executable, "-m", "pip", "install", 19                               "transformers==4.54.0",20                               "torchvision==0.22.1",21                               "diffusers==0.34.0",22                               "einops==0.8.1"])23 24ensure_env_installed()25 26# Global model variable27model_zoo = {28    "imuru_small": {29        "repo_id": "Ruian7P/imuru_small",30    },31    "imuru_large": {32        "repo_id": "Ruian7P/imuru_large",33    },34    # "emuru_t5_small": {35    #     "repo_id": "Ruian7P/emuru_result",36    #     "model_name": "emuru_t5_small_2e-5_ech5"37    # }38}39 40model = None41 42def load_model(model_name="imuru_large"):    43    global model44 45    if model is None:46        print(f"Loading model {model_name}...")47        from transformers import AutoModel48        49        model = AutoModel.from_pretrained(50            model_zoo[model_name]["repo_id"],51            trust_remote_code=True52        )53        model.eval()54        print("✅ Model loaded")55    56    return model57 58 59def load_examples():60    """Load example samples."""61    examples = []62    examples.append([63        "sample/sample.png", "Ruian7P"64    ])65    return examples66 67def process_image(img):68    from torchvision.transforms import functional as F69    img = img.convert("RGB")70    img = img.resize((img.width * 64 // img.height, 64))71    img = F.to_tensor(img)72    img = F.normalize(img, [0.5], [0.5])73    return img74 75 76@spaces.GPU77def generate_handwriting(style_image, gen_text, model_name="imuru_large"):78    """Generate handwriting in the style of the input image."""79    if not gen_text or gen_text.strip() == "":80        return None, "❌ Please provide text to generate"81    82    if style_image is None:83        return None, "❌ Please upload a style image"84    85    try:86        # Convert numpy array to PIL Image if needed87        if isinstance(style_image, np.ndarray):88            style_image = Image.fromarray(style_image)89        90        # Load and move model to GPU91        loaded_model = load_model(model_name)92        loaded_model.to("cuda")93        94        # Preprocess style image95        style_img = process_image(style_image).to("cuda")96        97        # Generate98        with torch.inference_mode():99            result = loaded_model.generate(100                style_img=style_img,101                gen_text=gen_text,102                max_new_tokens=512103            )104        105        return result, "✅ Generation successful!"106        107    except Exception as e:108        import traceback109        traceback.print_exc()110        return None, f"❌ Error: {str(e)}"111 112 113# Custom CSS for better styling114custom_css = """115.gradio-container {116    width: 100%;117    max-width: 1200px !important;118    margin: 0 auto !important;119}120.header-text {121    text-align: center;122    margin-bottom: 1rem;123}124.feature-box {125    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);126    border-radius: 10px;127    padding: 15px;128    margin: 10px 0;129}130footer {131    visibility: hidden;132}133"""134 135# Build the interface with gr.Blocks for better customization136with gr.Blocks(css=custom_css, title="Imuru") as demo:137    138    # Header139    gr.HTML("""140    <div style="text-align: center; margin-bottom: 20px;">141        <h1>🍎 Imuru: Autoregressive Handwriting Generation</h1>142    </div>143    """)144    145    with gr.Row():146        with gr.Column(scale=1):147            model_selector = gr.Dropdown(148                label="🤖 Select Model",149                choices=list(model_zoo.keys()),150                value="imuru_large",151                interactive=True152            )153 154            style_image_input = gr.Image(155                label="🖼️ Style Image",156                type="pil",157                height=200158            )159            160            gen_text_input = gr.Textbox(161                label="✍️ Text to Generate",162                placeholder="Enter the text you want to generate in the selected style",163                lines=2,164                value="Hello, I am Imuru!"165            )166            167            generate_btn = gr.Button("🕶 Generate", variant="primary", size="lg")168        169        with gr.Column(scale=1):170            output_image = gr.Image(171                label="🖼️ Generated Output",172                type="pil",173                height=200174            )175            176            status_text = gr.Textbox(177                label="🚧 Status",178                lines=1,179                interactive=False180            )181    182    # Examples183    examples = load_examples()184    if examples:185        gr.Examples(186            examples=examples,187            inputs=[style_image_input, gen_text_input],188            label="💡 Examples",189            examples_per_page=4190        )191    192    # Connect events193    generate_btn.click(194        fn=generate_handwriting,195        inputs=[style_image_input, gen_text_input, model_selector],196        outputs=[output_image, status_text]197    )198    199    gen_text_input.submit(200        fn=generate_handwriting,201        inputs=[style_image_input, gen_text_input, model_selector],202        outputs=[output_image, status_text]203    )204    205    # How to use section206    gr.Markdown("""207    ---208    ### 🧠 How to Use209    210    1. **Upload a style image**: A handwritten sample to extract style from211    2. **Type generation text**: The text you want to generate in the style of the image212    3. **Click Generate**: Imuru will create the handwritten text image for you!213    """)214 215 216if __name__ == "__main__":217    demo.launch()218