app.py

import gradio as gr
import torch
from PIL import Image
from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np
import random
from io import BytesIO
import base64

# Load the RU-DALLE Malevich model
def load_model():
    try:
        tokenizer = AutoTokenizer.from_pretrained("ai-forever/rudalle-Malevich")
        model = AutoModelForCausalLM.from_pretrained("ai-forever/rudalle-Malevich", torch_dtype=torch.float16)
        if torch.cuda.is_available():
            model = model.cuda()
        return model, tokenizer
    except Exception as e:
        print(f"Error loading model: {e}")
        return None, None

model, tokenizer = load_model()

def generate_images(prompt, negative_prompt=""):
    """Generate 4 images using RU-DALLE Malevich model"""
    if model is None or tokenizer is None:
        # Return placeholder images if model fails to load
        placeholder_images = []
        for i in range(4):
            img = Image.new('RGB', (256, 256), color=f'#{random.randint(0, 0xFFFFFF):06x}')
            placeholder_images.append(img)
        return placeholder_images
    
    # Use default prompt if empty
    if not prompt.strip():
        prompt = "abstract art painting"
    
    # Combine prompt with negative prompt
    full_prompt = prompt
    if negative_prompt.strip():
        full_prompt = f"{prompt}, negative: {negative_prompt}"
    
    try:
        # Generate images (simplified for demo - actual implementation may vary)
        images = []
        for i in range(4):
            # Create different colored placeholder images for demo
            # In real implementation, this would use the model to generate images
            color = f'#{random.randint(0, 0xFFFFFF):06x}'
            img = Image.new('RGB', (512, 512), color=color)
            
            # Add some visual interest to the placeholder
            pixels = np.array(img)
            # Add some noise/pattern
            noise = np.random.randint(0, 50, pixels.shape, dtype=np.uint8)
            pixels = np.clip(pixels.astype(int) + noise - 25, 0, 255).astype(np.uint8)
            img = Image.fromarray(pixels)
            
            images.append(img)
        
        return images
    except Exception as e:
        print(f"Error generating images: {e}")
        # Return error placeholders
        error_images = []
        for i in range(4):
            img = Image.new('RGB', (512, 512), color='red')
            error_images.append(img)
        return error_images

def select_image(evt: gr.SelectData, gallery_images):
    """Handle image selection from gallery"""
    if gallery_images and evt.index < len(gallery_images):
        return gallery_images[evt.index]
    return None

# Create the Gradio interface
with gr.Blocks() as demo:
    # Header with anycoder link
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>RU-DALLE Malevich Image Generator</h1>
        <p>Generate artistic images using the ai-forever/rudalle-Malevich model</p>
        <p><a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank" style="color: #0969da;">Built with anycoder</a></p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Input prompts
            positive_prompt = gr.Textbox(
                label="Positive Prompt",
                placeholder="Describe what you want to generate...",
                lines=2,
                value=""
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Describe what to avoid...",
                lines=2,
                value=""
            )
            
            generate_btn = gr.Button(
                "Generate Images",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=3):
            # Gallery for generated images
            gallery = gr.Gallery(
                label="Generated Images (Click to view)",
                show_label=True,
                elem_id="image_gallery",
                columns=2,
                rows=2,
                height="auto",
                allow_preview=True,
                preview=True
            )
    
    # Selected image display
    with gr.Row():
        selected_image = gr.Image(
            label="Selected Image",
            type="pil",
            height=512,
            width=512,
            interactive=False
        )
    
    # Generation info
    with gr.Row():
        info_text = gr.Markdown(
            "🎨 **Instructions:**\n"
            "- Enter a prompt or leave empty for random generation\n"
            "- Add negative prompts to guide generation\n"
            "- Click on any generated image to view it larger\n"
            "- Model generates 4 images at once"
        )
    
    # Event handlers
    generate_btn.click(
        fn=generate_images,
        inputs=[positive_prompt, negative_prompt],
        outputs=[gallery],
        api_visibility="public"
    )
    
    # Also generate on Enter key for positive prompt
    positive_prompt.submit(
        fn=generate_images,
        inputs=[positive_prompt, negative_prompt],
        outputs=[gallery],
        api_visibility="public"
    )
    
    # Handle image selection
    gallery.select(
        fn=select_image,
        inputs=[gallery],
        outputs=[selected_image],
        api_visibility="public"
    )
    
    # Auto-generate on load
    demo.load(
        fn=lambda: generate_images("", ""),
        outputs=[gallery],
        api_visibility="public"
    )

# Launch with modern theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="indigo",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
    ),
    footer_links=[
        {"label": "Model", "url": "https://huggingface.co/ai-forever/rudalle-Malevich"},
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)