app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from scratch3 import Project

# Load the LLaMa model and tokenizer
model_name = "TheBloke/LLaMa-7B-GGML"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Define extended Scratch blocks
scratch_blocks = {
    "motion": [
        {"name": "move", "parameters": {"steps": "number"}},
        {"name": "turn", "parameters": {"angle": "number"}},
        {"name": "go to", "parameters": {"x": "number", "y": "number"}},
        {"name": "point in direction", "parameters": {"direction": "number"}}
    ],
    "sound": [
        {"name": "play sound", "parameters": {"sound": "string"}},
        {"name": "stop all sounds", "parameters": {}}
    ],
    "event": [
        {"name": "when green flag clicked", "parameters": {}},
        {"name": "when this sprite clicked", "parameters": {}}
    ],
    "control": [
        {"name": "wait", "parameters": {"seconds": "number"}},
        {"name": "repeat", "parameters": {"times": "number"}},
        {"name": "if then", "parameters": {"condition": "string"}},
    ],
    "variables": [
        {"name": "set variable", "parameters": {"variable": "string", "value": "number"}},
        {"name": "change variable", "parameters": {"variable": "string", "value": "number"}},
    ],
    "looks": [
        {"name": "say", "parameters": {"message": "string"}},
        {"name": "think", "parameters": {"message": "string"}},
        {"name": "show", "parameters": {}},
        {"name": "hide", "parameters": {}}
    ],
    "drawing": [
        {"name": "draw line", "parameters": {"x1": "number", "y1": "number", "x2": "number", "y2": "number"}},
        {"name": "draw circle", "parameters": {"x": "number", "y": "number", "radius": "number"}},
        {"name": "draw rectangle", "parameters": {"x": "number", "y": "number", "width": "number", "height": "number"}},
    ]
}

def generate_script(prompt):
    """
    Generate a Scratch script based on the provided prompt using the AI model.
    """
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(inputs["input_ids"], max_length=300, num_return_sequences=1)
    
    generated_script = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_script

def parse_generated_script(generated_script):
    """
    Parse the generated script into instructions.
    """
    instructions = []
    for line in generated_script.split('.'):
        line = line.strip()
        if line:
            instructions.append(line)
    return instructions

def create_scratch_project(title, sprites, sounds, custom_prompt):
    """
    Create a Scratch project based on user input and generated script.
    """
    # Combine title and custom prompt for better context
    prompt = f"Create a Scratch project titled '{title}' featuring sprites {sprites} with sounds {sounds}. " \
             f"Additional details: {custom_prompt}"
    generated_script = generate_script(prompt)

    # Parse the generated script into instructions
    instructions = parse_generated_script(generated_script)

    # Create a new Scratch project
    project = Project()

    # Add sprites
    sprite_list = [sprite.strip() for sprite in sprites.split(',')]
    for sprite in sprite_list:
        sprite_block = {
            "objName": sprite,
            "skin": 1,
            "x": 0,
            "y": 0,
            "direction": 90,
            "visible": True,
            "size": 100
        }
        project.sprites.append(sprite_block)

    # Add sounds if provided
    for sound in sounds.split(','):
        sound = sound.strip()
        if sound:
            sound_block = {
                "name": sound,
                "path": f"sounds/{sound}.wav"
            }
            project.sounds.append(sound_block)

    # Map parsed instructions to Scratch blocks
    for instruction in instructions:
        if "move" in instruction:
            steps = extract_number(instruction)
            project.add_motion_block({"name": "move", "parameters": {"steps": steps}})
        elif "turn" in instruction:
            angle = extract_number(instruction)
            project.add_motion_block({"name": "turn", "parameters": {"angle": angle}})
        elif "play sound" in instruction:
            project.add_sound_block({"name": "play sound", "parameters": {"sound": sounds.split(',')[0].strip()}})
        elif "when green flag clicked" in instruction:
            project.add_event_block({"name": "when green flag clicked"})
        elif "repeat" in instruction:
            times = extract_number(instruction)
            project.add_control_block({"name": "repeat", "parameters": {"times": times}})
        elif "say" in instruction:
            message = extract_message(instruction)
            project.add_looks_block({"name": "say", "parameters": {"message": message}})
        elif "draw line" in instruction:
            x1, y1, x2, y2 = extract_drawing_parameters(instruction)
            project.add_drawing_block({"name": "draw line", "parameters": {"x1": x1, "y1": y1, "x2": x2, "y2": y2}})
        elif "draw circle" in instruction:
            x, y, radius = extract_circle_parameters(instruction)
            project.add_drawing_block({"name": "draw circle", "parameters": {"x": x, "y": y, "radius": radius}})
        elif "draw rectangle" in instruction:
            x, y, width, height = extract_rectangle_parameters(instruction)
            project.add_drawing_block({"name": "draw rectangle", "parameters": {"x": x, "y": y, "width": width, "height": height}})
        elif "remove sprite" in instruction:
            sprite_name = extract_sprite_name(instruction)
            project.remove_sprite(sprite_name)
        elif "rename sprite" in instruction:
            old_name, new_name = extract_rename_sprite_parameters(instruction)
            project.rename_sprite(old_name, new_name)

    # Save the project as an .sb3 file
    sb3_file_path = f"{title}.sb3"
    with open(sb3_file_path, "wb") as f:
        f.write(project.export())
    
    return sb3_file_path

def extract_number(instruction):
    """
    Extract a number from the instruction string.
    """
    import re
    match = re.search(r'\d+', instruction)
    return int(match.group(0)) if match else 0

def extract_message(instruction):
    """
    Extract a message from the instruction string.
    """
    return instruction.split("say")[-1].strip().strip('"')

def extract_drawing_parameters(instruction):
    """
    Extract parameters for drawing a line from the instruction string.
    """
    numbers = list(map(int, re.findall(r'\d+', instruction)))
    return numbers

def extract_circle_parameters(instruction):
    """
    Extract parameters for drawing a circle from the instruction string.
    """
    numbers = list(map(int, re.findall(r'\d+', instruction)))
    return numbers

def extract_rectangle_parameters(instruction):
    """
    Extract parameters for drawing a rectangle from the instruction string.
    """
    numbers = list(map(int, re.findall(r'\d+', instruction)))
    return numbers

def extract_sprite_name(instruction):
    """
    Extract the sprite name from the instruction string for removal.
    """
    return instruction.split("remove sprite")[-1].strip()

def extract_rename_sprite_parameters(instruction):
    """
    Extract old and new names for renaming a sprite from the instruction string.
    """
    parts = instruction.split("rename sprite")[-1].strip().split("to")
    old_name = parts[0].strip()
    new_name = parts[1].strip() if len(parts) > 1 else None
    return old_name, new_name

# Create Gradio interface
iface = gr.Interface(
    fn=create_scratch_project,
    inputs=[
        gr.Textbox(label="Project Title"),
        gr.Textbox(label="Sprite Names (comma separated)"),
        gr.Textbox(label="Sound Names (comma separated)"),
        gr.Textbox(label="Custom Prompt (optional)"),
    ],
    outputs=gr.File(label="Download Your Scratch Project (.sb3)"),
    title="Scratch Project Generator",
    description="Create a Scratch project using AI. Provide a title, sprites, sounds, and a custom prompt."
)

# Launch the Gradio app
iface.launch()