# Project export: AI Art Alchemist: Reinforcement Learning Guided Diffusion

This document was generated by HackStack to give an AI agent context about a hackathon project. Sections are labeled with their provenance; content marked as truncated was cut to keep this document small.

## Project metadata

- Hackathon: UC Berkeley AI Hackathon 2025
- Tagline: We use an RL agent to dynamically mix specialized AI art diffusion models at each step, creating images superior to what any single model can achieve alone.
- Devpost: https://devpost.com/software/reinforcement-learning-guided-diffusion
- GitHub: https://github.com/jaisharmz/berkeleyaihackathon2025
- Demo: https://docs.google.com/presentation/d/1IQQjCY09wp_LcUZ086dxvZsmYHhpeRpGBJiHxTVbQRY/edit?usp=sharing
- Team: 1 GitHub contributor(s) — jaisharmz (4 commits)

## Devpost submission (written by the team)

No Devpost description available.

## README (from the GitHub repository)

# berkeleyaihackathon2025
second project for berkeley ai hackathon 2025


## Detected evidence (automated analysis)

Indexed codebase: 11 recognized source files, 71 KB.
- Python (language) — detected in the code
- Hugging Face (technology) — claimed on Devpost, not found in the code

## Codebase structure (from repository index)

### Files (15 of 15)

```
code_gen_output/code_fixed_workflow.py
code_gen_output/code_logic_only.py
code_gen_output/code_rl_guided.py
codewriting1.py
inpainting_output/choice_history.txt
main.py
main2.py
main3.py
main4.py
main5.py
README.md
story_output/story_dialogue_only.txt
story_output/story_plot_only.txt
story_output/story_rl_guided.txt
storywriting1.py
```

### Dependencies

No dependency index available.

### Recent commits (newest first)

- Merge branch 'main' of https://github.com/jaisharmz/berkeleyaihackathon2025
- fixing
- Initial commit
- Initial commit

## Key source files (fetched from GitHub, selected and truncated for size)

### main.py

```python
import torch
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from diffusers import DDIMScheduler, StableDiffusionPipeline
import numpy as np
import os
from tqdm.auto import tqdm

# --- Configuration ---
# --------------------------------------------------------------------------------------
# Model Identifiers from Hugging Face Hub
MODEL_A_ID = "runwayml/stable-diffusion-v1-5"
MODEL_B_ID = "prompthero/openjourney" # A Midjourney-style fine-tune of SD 1.5

# Reinforcement Learning Agent (Multi-Armed Bandit) Settings
EPSILON = 0.3  # 30% chance to explore, 70% chance to exploit
# --- NEW PROMPT FOR EXPERIMENTATION ---
# Try other prompts here to see how the agent behaves!
# PROMPT = "A blueprint schematic of a fantasy dragon."
PROMPT = "A photograph of a New York City street, reimagined by Van Gogh."
# PROMPT = "A photorealistic portrait of an astronaut, in the style of a detailed oil painting."

# Diffusion Process Settings
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 7.5
HEIGHT = 512
WIDTH = 512
SEED = 42

# Output settings
OUTPUT_DIR = "rl_diffusion_output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------------------------------------------


class MultiArmedBandit:
    """
    A simple Epsilon-Greedy Multi-Armed Bandit (MAB) agent.
    This agent will decide which diffusion model to use at each step.
    """
    def __init__(self, model_names, epsilon):
        self.models = {name: {'pulls': 0, 'value': 0.0} for name in model_names}
        self.epsilon = epsilon
        self.names = model_names

    def select_action(self):
        """
        Chooses a model (arm) to pull and increments its pull count.
        """
        # --- FIX: Increment pull count immediately upon selection ---
        if np.random.uniform(0, 1) < self.epsilon:
            # Explore: choose a random model
            action = np.random.choice(self.names)
            print(f"  >> Agent decision: EXPLORE -> Chose {action}")
        else:
            # Exploit: choose the model with the current highest value
            action = max(self.models, key=lambda m: self.models[m]['value'])
            print(f"  >> Agent decision: EXPLOIT -> Chose {action} (Best value: {self.models[action]['value']:.4f})")
        
        # Increment the pull count for the chosen action
        self.models[action]['pulls'] += 1
        return action

    def update_value(self, model_name, reward):
        """
        Updates the value of a model based on the received reward.
        The pull count is NOT changed here.
        """
        model_stats = self.models[model_name]
        # Use an incremental average formula to update the value
        # Note: We use the current number of pulls for the averaging
        model_stats['value'] = model_stats['value'] + (1 / model_stats['pulls']) * (reward - model_stats['value'])
        print(f"  >> Agent updated: {model_name} now has value {model_stats['value']:.4f} after {model_stats['pulls']} pulls.")

class RLDrivenDiffusion:
    """
    Manages the diffusion process guided by the RL agent.
    """
    def __init__(self, prompt, model_a_id, model_b_id, seed):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")

        self.prompt = prompt
        self.seed = seed

        # --- Load Models ---
        print("Loading models... This may take a while.")
        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(self.device)
        self.pipe_a = StableDiffusionPipeline.from_pretrained(model_a_id, torch_dtype=torch.float16).to(self.device)
        self.pipe_b = StableDiffusionPipeline.from_pretrained(model_b_id, torch_dtype=torch.float16).to(self.device)
        self.scheduler = DDIMScheduler.from_config(self.pipe_a.scheduler.config)

        self.tokenizer = self.pipe_a.tokenizer
        self.text_encoder = self.pipe_a.text_encoder
        self.unet_a = self.pipe_a.unet
        self.unet_b = self.pipe_b.unet
        self.vae = self.pipe_a.vae
        
        print("Models loaded successfully.")

    def get_clip_score(self, image):
        with torch.no_grad():
            inputs = self.clip_processor(text=[self.prompt], images=[image], return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            outputs = self.clip_model(**inputs)
            return outputs.logits_per_image.item()

    def _get_text_embeddings(self, prompt):
        text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
        with torch.no_grad():
            text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
        return text_embeddings

    def generate(self, agent=None, force_model=None, output_filename="final_image.png"):
        """
        The main generation loop.
        - If 'agent' is provided, runs in RL-guided mode.
        - If 'force_model' is 'Model_A' or 'Model_B', uses only that model.
        """
        # --- 1. Setup ---
        uncond_embeddings = self._get_text_embeddings("")
        text_embeddings = self._get_text_embeddings(self.prompt)
        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
        
        self.scheduler.set_timesteps(NUM_INFERENCE_STEPS, device=self.device)
        
        generator = torch.Generator(device=self.device).manual_seed(self.seed)
        latents = torch.randn((1, self.unet_a.config.in_channels, HEIGHT // 8, WIDTH // 8), generator=generator, device=self.device, dtype=torch.float16)
        latents = latents * self.scheduler.init_noise_sigma
        
        last_score = 0.0
        
        # --- 2. The Diffusion Loop ---
        for i, t in enumerate(tqdm(self.scheduler.timesteps, desc=f"Generating with {force_model or 'RL Agent'}")):
  
[truncated — 4141 more characters]
```

### storywriting1.py

```python
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration
from sentence_transformers import SentenceTransformer, util
import math
import os
from tqdm.auto import tqdm
import re

# --- Configuration ---
# --------------------------------------------------------------------------------------
# User Prompt and its breakdown for the reward function
USER_PROMPT = "A futuristic sci-fi noir mystery with witty back-and-forth dialogue."
# These phrases will be embedded to guide the different models
PROMPT_ASPECTS = {
    "Plot": "A futuristic sci-fi noir mystery with suspense and action.",
    "Scene": "A detailed and atmospheric description of a futuristic city.",
    "Dialogue": "Witty, sharp, back-and-forth dialogue between characters."
}

# Reinforcement Learning Agent (UCB) Settings
UCB_EXPLORATION_CONSTANT = 1.5  # Balances exploration and exploitation

# Story Generation Settings
NUM_SENTENCES = 15  # The total number of sentences in the final story
SEED = 42

# Output settings
OUTPUT_DIR = "story_output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------------------------------------------

class UCB1_Bandit:
    """The UCB1 Multi-Armed Bandit agent - our 'Head Writer'."""
    def __init__(self, model_names, exploration_constant):
        self.models = {name: {'pulls': 0, 'value': 0.0} for name in model_names}
        self.exploration_constant = exploration_constant
        self.total_pulls = 0
        self.names = model_names

    def select_action(self):
        self.total_pulls += 1
        for name in self.names:
            if self.models[name]['pulls'] == 0:
                print(f"  >> Agent Decision: Initializing -> Chose {name}")
                self.models[name]['pulls'] += 1
                return name
        
        ucb_scores = {}
        for name in self.names:
            model_stats = self.models[name]
            average_reward = model_stats['value']
            exploration_bonus = self.exploration_constant * math.sqrt(
                math.log(self.total_pulls) / model_stats['pulls']
            )
            ucb_scores[name] = average_reward + exploration_bonus
        
        action = max(ucb_scores, key=ucb_scores.get)
        self.models[action]['pulls'] += 1
        print(f"  >> Agent Decision: UCB -> Chose {action}")
        return action

    def update_value(self, model_name, reward):
        model_stats = self.models[model_name]
        model_stats['value'] = ((model_stats['value'] * (model_stats['pulls'] - 1)) + reward) / model_stats['pulls']
        print(f"  >> Agent Update: {model_name} value is now {model_stats['value']:.3f}")

class AI_Writers_Room:
    """Manages the models, generation process, and reward calculation."""
    def __init__(self, prompt_aspects):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")

        # --- Load Models ---
        print("Loading models from Hugging Face Hub...")
        self.reward_model = SentenceTransformer('all-MiniLM-L6-v2', device=self.device)
        self.tokenizer = T5Tokenizer.from_pretrained('google-t5/t5-small')
        self.model = T5ForConditionalGeneration.from_pretrained('google-t5/t5-small').to(self.device)
        print("Models loaded successfully.")

        self.prompt_embeddings = {}
        for aspect, text in prompt_aspects.items():
            self.prompt_embeddings[aspect] = self.reward_model.encode(text, convert_to_tensor=True)
            
        self.world_state = {
            "characters": ["Detective Kaito", "a mysterious informant named Anya"],
            "setting": "a rain-slicked street in Neo-Kyoto under a perpetual twilight sky",
            "last_speaker": None,
            "initial_story": "Detective Kaito pulled the collar of his trench coat tighter, the neon signs of Neo-Kyoto reflecting in the puddles at his feet."
        }

    def _generate_with_t5(self, input_text):
        """Helper for T5-based models with improved generation parameters."""
        inputs = self.tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True).to(self.device)
        outputs = self.model.generate(
            inputs.input_ids, 
            max_new_tokens=60, 
            num_beams=5, 
            no_repeat_ngram_size=2, 
            early_stopping=True,
            temperature=0.9
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

    # --- The "Specialized Writers" ---
    def plot_driver(self, story_context):
        # FINAL FIX: Use a simple, unambiguous instruction prefix.
        prompt = f"Write a single sentence to continue the following story: {story_context}"
        return self._generate_with_t5(prompt)

    def scene_setter(self, story_context):
        # FINAL FIX: Use a simple, unambiguous instruction prefix.
        prompt = f"Write a single sentence to describe the setting of the following story in more detail: {story_context}"
        return self._generate_with_t5(prompt)

    def dialogue_writer(self, story_context):
        # FINAL FIX: Use a simple, unambiguous instruction prefix for dialogue.
        if self.world_state['last_speaker'] == "Detective Kaito":
            next_speaker = self.world_state['characters'][1] # Anya
        else:
            next_speaker = "Detective Kaito"
        
        last_line = re.split(r'(?<=[.!?"])\s+', story_context.strip())[-1]
        prompt = (f"The last line of dialogue was: {last_line}. "
                  f"Write the next line of dialogue spoken by {next_speaker}.")
        response = self._generate_with_t5(prompt)
        self.world_state['last_speaker'] = next_speaker
        return f'"{response}"'

    def calculate_reward(self, generated_sentence, model_name):
        """Calculates reward based on semantic similarity to the prompt aspect."""
        if not generated_sentence or not generated_sentence.strip():
            return 0.0
        sentence_embeddi
[truncated — 3030 more characters]
```

### main4.py

```python
import torch
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel, CLIPProcessor
from PIL import Image
import os
from tqdm.auto import tqdm
import numpy as np

# --- Configuration ---
# Models
MODEL_A_ID = "runwayml/stable-diffusion-v1-5"
MODEL_B_ID = "runwayml/stable-diffusion-inpainting"
CLIP_ID = "openai/clip-vit-large-patch14"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32

# Generation Parameters
PROMPT = "a majestic fantasy castle in a lush valley, high quality, 4k, digital art"
NUM_INFERENCE_STEPS = 50  # Total number of diffusion steps
GUIDANCE_SCALE = 8.0      # How much to adhere to the prompt

# RL / Bandit Parameters
DECISION_INTERVAL = 5 # How many steps to run a model before re-evaluating
EPSILON = 0.3         # Probability of choosing a random model (exploration)
OUTPUT_DIR = "inpainting_output"

def decode_latents_to_pil(vae, latents, dtype):
    """Helper function to decode latents into a PIL image."""
    with torch.no_grad():
        # Move latents to CPU for decoding if they are on CUDA
        latents_for_decode = latents.to(dtype)
        image = vae.decode(1 / 0.18215 * latents_for_decode).sample
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
        return Image.fromarray((image[0] * 255).round().astype(np.uint8))

def main():
    """
    Main function to run the multi-model generation process and individual model benchmarks.
    """
    print(f"Using device: {DEVICE}")

    # --- 1. Load Models ---
    print("Loading models...")
    tokenizer = CLIPTokenizer.from_pretrained(MODEL_A_ID, subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained(MODEL_A_ID, subfolder="text_encoder", torch_dtype=DTYPE).to(DEVICE)
    vae = AutoencoderKL.from_pretrained(MODEL_A_ID, subfolder="vae", torch_dtype=DTYPE).to(DEVICE)
    
    # Model A & B
    unet_A = UNet2DConditionModel.from_pretrained(MODEL_A_ID, subfolder="unet", torch_dtype=DTYPE).to(DEVICE)
    unet_B = UNet2DConditionModel.from_pretrained(MODEL_B_ID, subfolder="unet", torch_dtype=DTYPE).to(DEVICE)

    scheduler = PNDMScheduler.from_pretrained(MODEL_A_ID, subfolder="scheduler")

    # Reward Model
    clip_processor = CLIPProcessor.from_pretrained(CLIP_ID)
    clip_model = CLIPModel.from_pretrained(CLIP_ID, torch_dtype=DTYPE).to(DEVICE)
    print("Models loaded.")

    # --- 2. Helper function for CLIP score ---
    def calculate_clip_score(image, text):
        with torch.no_grad():
            inputs = clip_processor(text=[text], images=[image], return_tensors="pt", padding=True)
            inputs = {k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k,v in inputs.items()}
            if DTYPE == torch.float16:
                inputs['pixel_values'] = inputs['pixel_values'].to(DTYPE)
            outputs = clip_model(**inputs)
            return outputs.logits_per_image.item()

    # --- 3. Prepare for Generation ---
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    text_input = tokenizer(PROMPT, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
    with torch.no_grad():
        text_embeddings = text_encoder(text_input.input_ids.to(DEVICE))[0].to(DTYPE)
    uncond_input = tokenizer([""], padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt")
    with torch.no_grad():
        uncond_embeddings = text_encoder(uncond_input.input_ids.to(DEVICE))[0].to(DTYPE)
    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

    height, width = 512, 512
    generator = torch.manual_seed(42)
    initial_latents = torch.randn(
        (1, unet_A.config.in_channels, height // 8, width // 8),
        generator=generator,
        device='cpu'
    ).to(DEVICE, dtype=DTYPE)

    scheduler.set_timesteps(NUM_INFERENCE_STEPS)

    # --- 4. Mixed-Model Generation (RL Agent with IMPROVED REWARDS) ---
    print("\n--- Starting Mixed-Model Generation (RL Agent with Improved Rewards) ---")
    latents = initial_latents.clone()
    models = {"A": unet_A, "B": unet_B}
    model_names = list(models.keys())
    model_rewards = {name: 0.0 for name in model_names}
    model_counts = {name: 0 for name in model_names}
    choice_history = []
    
    # **MODIFICATION**: Initialize last_clip_score to track score changes
    last_clip_score = 0.0

    for i, t in enumerate(tqdm(scheduler.timesteps, desc="Mixed Model")):
        if i % DECISION_INTERVAL == 0:
            if torch.rand(1).item() < EPSILON or all(c == 0 for c in model_counts.values()):
                chosen_model_name = model_names[torch.randint(0, len(model_names), (1,)).item()]
            else:
                avg_rewards = {name: model_rewards[name] / model_counts[name] for name in model_names if model_counts[name] > 0}
                chosen_model_name = max(avg_rewards, key=avg_rewards.get)
        
        choice_history.append(chosen_model_name)
        unet = models[chosen_model_name]

        if chosen_model_name == 'B':
            mask = torch.zeros_like(latents[:, :1])
            masked_image_latents = torch.zeros_like(latents)
            nine_channel_latents = torch.cat([latents, mask, masked_image_latents], dim=1)
            latent_model_input = torch.cat([nine_channel_latents] * 2)
        else:
            latent_model_input = torch.cat([latents] * 2)

        latent_model_input = scheduler.scale_model_input(latent_model_input, t)
        with torch.no_grad():
            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
        
        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
        noise_pred = noise_pred_uncond + GUIDANCE_SCALE * (noise_pred_text - noise_pred_uncond)
        latents = scheduler.step(noise_pred, t, latents).prev_sample

        if (i + 1) % DECISION_INTERVAL == 0 and
[truncated — 3262 more characters]
```

### codewriting1.py

```python
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
from sentence_transformers import SentenceTransformer, util
import math
import os
import ast
from tqdm.auto import tqdm
import re

# --- Configuration ---
# --------------------------------------------------------------------------------------
USER_PROMPT = "Write a Python function that takes a list of strings and returns a new list with only the strings that are palindromes."

# Reinforcement Learning Agent (UCB) Settings
UCB_EXPLORATION_CONSTANT = 2.5

# Code Generation Settings
NUM_BLOCKS_TO_GENERATE = 15
SEED = 42

# Model and Output settings
# FIX: Upgraded to a larger, more capable model for better results.
CODE_MODEL_ID = "Salesforce/codet5-large"
REWARD_MODEL_ID = "all-MiniLM-L6-v2"
OUTPUT_DIR = "code_gen_output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------------------------------------------

class UCB1_Bandit:
    """The UCB1 Multi-Armed Bandit agent - our 'Tech Lead'."""
    def __init__(self, model_names, exploration_constant):
        self.models = {name: {'pulls': 0, 'value': 0.0} for name in model_names}
        self.exploration_constant = exploration_constant
        self.total_pulls = 0
        self.names = model_names

    def select_action(self):
        self.total_pulls += 1
        for name in self.names:
            if self.models[name]['pulls'] == 0:
                print(f"  >> Agent Decision: Initializing -> Chose {name}")
                self.models[name]['pulls'] += 1
                return name
        
        ucb_scores = {}
        for name in self.names:
            model_stats = self.models[name]
            average_reward = model_stats['value']
            exploration_bonus = self.exploration_constant * math.sqrt(
                math.log(self.total_pulls) / self.models[name]['pulls']
            )
            ucb_scores[name] = average_reward + exploration_bonus
        
        action = max(ucb_scores, key=ucb_scores.get)
        self.models[action]['pulls'] += 1
        print(f"  >> Agent Decision: UCB -> Chose {action}")
        return action

    def update_value(self, model_name, reward):
        model_stats = self.models[model_name]
        model_stats['value'] = ((model_stats['value'] * (model_stats['pulls'] - 1)) + reward) / model_stats['pulls']
        print(f"  >> Agent Update: {model_name} value is now {model_stats['value']:.3f} after a reward of {reward:.2f}")

class AI_Pair_Programmer:
    """Manages the models, code generation process, and reward calculation."""
    def __init__(self, user_prompt):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")
        self.user_prompt = user_prompt

        print("Loading models...")
        self.reward_model = SentenceTransformer(REWARD_MODEL_ID, device=self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(CODE_MODEL_ID)
        self.model = T5ForConditionalGeneration.from_pretrained(CODE_MODEL_ID).to(self.device)
        print("Models loaded.")

        self.prompt_embedding = self.reward_model.encode(user_prompt, convert_to_tensor=True)

    def _generate_code(self, prompt, max_tokens=120):
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True).to(self.device)
        outputs = self.model.generate(
            inputs.input_ids, max_new_tokens=max_tokens, num_beams=5, early_stopping=True, no_repeat_ngram_size=2
        )
        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)

    # --- The "Specialized Developers" ---
    def logic_scripter(self, function_signature, docstring, function_body):
        current_code = f"def {function_signature}:\n{docstring}"
        prompt = (f"Based on the following docstring, complete the Python code for the function body.\n\n"
                  f"Function:\n{current_code}\n\n    # Write the code here")
        generated_code = self._generate_code(prompt)
        body_match = re.search(r'def[^{]*?:\s*?"""[^"]*?"""\s*(.*)', generated_code, re.DOTALL)
        return body_match.group(1).strip() if body_match else generated_code


    def docstring_documenter(self, function_signature):
        prompt = f"Write a standard Python docstring for the function `def {function_signature}:` that performs the following task: '{self.user_prompt}'"
        docstring = self._generate_code(prompt, max_tokens=80)
        return f'    """{docstring}"""'

    def example_generator(self, full_function_code):
        prompt = f"Write one or more Python assert statements to test that the following function works correctly:\n\n{full_function_code}"
        return self._generate_code(prompt, max_tokens=60)

    # --- The "Code Reviewer" ---
    def calculate_reward(self, new_chunk, full_code, specialist_name):
        if not new_chunk or not new_chunk.strip(): return -1.0 

        try:
            ast.parse(full_code)
            syntax_reward = 0.5
        except (SyntaxError, IndentationError):
            return -2.0 

        specialist_reward = 0
        if specialist_name == "Docstring":
            clean_chunk = new_chunk.replace('"""', '').strip()
            doc_embedding = self.reward_model.encode(clean_chunk, convert_to_tensor=True)
            specialist_reward = util.pytorch_cos_sim(doc_embedding, self.prompt_embedding).item()
        elif specialist_name == "Example":
            try:
                exec(full_code)
                specialist_reward = 2.0 
            except AssertionError:
                specialist_reward = -2.0
            except Exception:
                specialist_reward = -1.0 
        
        return syntax_reward + specialist_reward

    def generate_function(self, agent=None, fixed_workflow=None, output_filename="code.py"):
        code_parts = {
            "signature": "find_palindromes(strings: list) -> list:",
            "docstring": "",
            "body": "",
            "examples": 
[truncated — 3546 more characters]
```

### main2.py

```python
import torch
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from diffusers import DDIMScheduler, StableDiffusionPipeline
import numpy as np
import os
import math
from tqdm.auto import tqdm

# --- Configuration ---
# --------------------------------------------------------------------------------------
# Model Identifiers from Hugging Face Hub
MODEL_A_ID = "runwayml/stable-diffusion-v1-5"
MODEL_B_ID = "prompthero/openjourney" # A Midjourney-style fine-tune of SD 1.5

# Reinforcement Learning Agent (UCB) Settings
# UCB is used instead of Epsilon-Greedy to encourage better exploration.
UCB_EXPLORATION_CONSTANT = 2.0  # Constant 'c' for the UCB1 formula. Higher values encourage more exploration.
# PROMPT = "A photorealistic portrait of an astronaut, in the style of a detailed oil painting."
# PROMPT = "A blueprint schematic of a fantasy dragon."
PROMPT = "A photograph of a New York City street, reimagined by Van Gogh."

# Diffusion Process Settings
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 7.5
HEIGHT = 512
WIDTH = 512
SEED = 42

# Output settings
OUTPUT_DIR = "rl_diffusion_output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------------------------------------------


class UCB1_Bandit:
    """
    An Upper Confidence Bound (UCB1) Multi-Armed Bandit agent.
    This provides a more robust exploration strategy than Epsilon-Greedy.
    """
    def __init__(self, model_names, exploration_constant):
        self.models = {name: {'pulls': 0, 'value': 0.0} for name in model_names}
        self.exploration_constant = exploration_constant
        self.total_pulls = 0
        self.names = model_names

    def select_action(self):
        """
        Chooses a model (arm) to pull using the UCB1 formula.
        The formula balances exploitation (highest average reward) and
        exploration (uncertainty about under-sampled arms).
        """
        self.total_pulls += 1

        # First, play each arm once to initialize
        for name in self.names:
            if self.models[name]['pulls'] == 0:
                print(f"  >> Agent decision: INITIALIZING -> Chose {name}")
                self.models[name]['pulls'] += 1
                return name
        
        # --- UCB1 Calculation ---
        ucb_scores = {}
        for name in self.names:
            model_stats = self.models[name]
            average_reward = model_stats['value']
            
            # The exploration term encourages trying less-pulled arms
            exploration_bonus = self.exploration_constant * math.sqrt(
                math.log(self.total_pulls) / model_stats['pulls']
            )
            
            ucb_scores[name] = average_reward + exploration_bonus
            print(f"  >> UCB Score for {name}: {average_reward:.3f} (avg) + {exploration_bonus:.3f} (bonus) = {ucb_scores[name]:.3f}")

        # Choose the arm with the highest UCB score
        action = max(ucb_scores, key=ucb_scores.get)
        self.models[action]['pulls'] += 1
        print(f"  >> Agent decision: UCB -> Chose {action}")
        return action

    def update_value(self, model_name, reward):
        """
        Updates the value of a model based on the received reward.
        """
        model_stats = self.models[model_name]
        # Use an incremental average formula to update the value
        # This is equivalent to sum_of_rewards / num_pulls
        model_stats['value'] = ((model_stats['value'] * (model_stats['pulls'] - 1)) + reward) / model_stats['pulls']
        print(f"  >> Agent updated: {model_name} now has value {model_stats['value']:.4f} after {model_stats['pulls']} pulls.")


class RLDrivenDiffusion:
    """
    Manages the diffusion process guided by the RL agent.
    """
    def __init__(self, prompt, model_a_id, model_b_id, seed):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")

        self.prompt = prompt
        self.seed = seed

        # --- Load Models ---
        print("Loading models... This may take a while.")
        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(self.device)
        self.pipe_a = StableDiffusionPipeline.from_pretrained(model_a_id, torch_dtype=torch.float16).to(self.device)
        self.pipe_b = StableDiffusionPipeline.from_pretrained(model_b_id, torch_dtype=torch.float16).to(self.device)
        self.scheduler = DDIMScheduler.from_config(self.pipe_a.scheduler.config)

        self.tokenizer = self.pipe_a.tokenizer
        self.text_encoder = self.pipe_a.text_encoder
        self.unet_a = self.pipe_a.unet
        self.unet_b = self.pipe_b.unet
        self.vae = self.pipe_a.vae
        
        print("Models loaded successfully.")

    def get_clip_score(self, image):
        with torch.no_grad():
            inputs = self.clip_processor(text=[self.prompt], images=[image], return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            outputs = self.clip_model(**inputs)
            return outputs.logits_per_image.item()

    def _get_text_embeddings(self, prompt):
        text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
        with torch.no_grad():
            text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
        return text_embeddings

    def generate(self, agent=None, force_model=None, output_filename="final_image.png"):
        """
        The main generation loop.
        - If 'agent' is provided, runs in RL-guided mode.
        - If 'force_model' is 'Model_A' or 'Model_B', uses only that model.
        """
        # --- 1. Setup ---
        uncond_embeddings = self._get_text_embeddings("")
        text_embeddings = self._get_text_embeddings(self.prompt)
        
[truncated — 4555 more characters]
```

### main3.py

```python
import torch
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from diffusers import DDIMScheduler, StableDiffusionPipeline
import numpy as np
import os
import math
from tqdm.auto import tqdm

# --- Configuration ---
# --------------------------------------------------------------------------------------
# Model Identifiers from Hugging Face Hub
MODEL_A_ID = "runwayml/stable-diffusion-v1-5"
MODEL_B_ID = "prompthero/openjourney" # A Midjourney-style fine-tune of SD 1.5

# Reinforcement Learning Agent (UCB) Settings
# UCB is used instead of Epsilon-Greedy to encourage better exploration.
# A higher constant encourages more exploration to achieve a more balanced pull distribution.
UCB_EXPLORATION_CONSTANT = 5.0  # Constant 'c' for the UCB1 formula. Higher values encourage more exploration.
PROMPT = "A photorealistic portrait of an astronaut, in the style of a detailed oil painting."
# PROMPT = "A blueprint schematic of a fantasy dragon."
# PROMPT = "A photograph of a New York City street, reimagined by Van Gogh."

# Diffusion Process Settings
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 7.5
HEIGHT = 512
WIDTH = 512
SEED = 42

# Output settings
OUTPUT_DIR = "rl_diffusion_output"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# --------------------------------------------------------------------------------------


class UCB1_Bandit:
    """
    An Upper Confidence Bound (UCB1) Multi-Armed Bandit agent.
    This provides a more robust exploration strategy than Epsilon-Greedy.
    """
    def __init__(self, model_names, exploration_constant):
        self.models = {name: {'pulls': 0, 'value': 0.0} for name in model_names}
        self.exploration_constant = exploration_constant
        self.total_pulls = 0
        self.names = model_names

    def select_action(self):
        """
        Chooses a model (arm) to pull using the UCB1 formula.
        The formula balances exploitation (highest average reward) and
        exploration (uncertainty about under-sampled arms).
        """
        self.total_pulls += 1

        # First, play each arm once to initialize
        for name in self.names:
            if self.models[name]['pulls'] == 0:
                print(f"  >> Agent decision: INITIALIZING -> Chose {name}")
                self.models[name]['pulls'] += 1
                return name
        
        # --- UCB1 Calculation ---
        ucb_scores = {}
        for name in self.names:
            model_stats = self.models[name]
            average_reward = model_stats['value']
            
            # The exploration term encourages trying less-pulled arms
            exploration_bonus = self.exploration_constant * math.sqrt(
                math.log(self.total_pulls) / model_stats['pulls']
            )
            
            ucb_scores[name] = average_reward + exploration_bonus
            print(f"  >> UCB Score for {name}: {average_reward:.3f} (avg) + {exploration_bonus:.3f} (bonus) = {ucb_scores[name]:.3f}")

        # Choose the arm with the highest UCB score
        action = max(ucb_scores, key=ucb_scores.get)
        self.models[action]['pulls'] += 1
        print(f"  >> Agent decision: UCB -> Chose {action}")
        return action

    def update_value(self, model_name, reward):
        """
        Updates the value of a model based on the received reward.
        """
        model_stats = self.models[model_name]
        # Use an incremental average formula to update the value
        # This is equivalent to sum_of_rewards / num_pulls
        model_stats['value'] = ((model_stats['value'] * (model_stats['pulls'] - 1)) + reward) / model_stats['pulls']
        print(f"  >> Agent updated: {model_name} now has value {model_stats['value']:.4f} after {model_stats['pulls']} pulls.")


class RLDrivenDiffusion:
    """
    Manages the diffusion process guided by the RL agent.
    """
    def __init__(self, prompt, model_a_id, model_b_id, seed):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Using device: {self.device}")

        self.prompt = prompt
        self.seed = seed

        # --- Load Models ---
        print("Loading models... This may take a while.")
        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(self.device)
        self.pipe_a = StableDiffusionPipeline.from_pretrained(model_a_id, torch_dtype=torch.float16).to(self.device)
        self.pipe_b = StableDiffusionPipeline.from_pretrained(model_b_id, torch_dtype=torch.float16).to(self.device)
        self.scheduler = DDIMScheduler.from_config(self.pipe_a.scheduler.config)

        self.tokenizer = self.pipe_a.tokenizer
        self.text_encoder = self.pipe_a.text_encoder
        self.unet_a = self.pipe_a.unet
        self.unet_b = self.pipe_b.unet
        self.vae = self.pipe_a.vae
        
        print("Models loaded successfully.")

    def get_clip_score(self, image):
        with torch.no_grad():
            inputs = self.clip_processor(text=[self.prompt], images=[image], return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            outputs = self.clip_model(**inputs)
            return outputs.logits_per_image.item()

    def _get_text_embeddings(self, prompt):
        text_input = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
        with torch.no_grad():
            text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
        return text_embeddings

    def generate(self, agent=None, force_model=None, output_filename="final_image.png"):
        """
        The main generation loop.
        - If 'agent' is provided, runs in RL-guided mode.
        - If 'force_model' is 'Model_A' or 'Model_B', uses only that model.
        """
        # --- 1. Setup ---
        uncond_embeddings = self._get
[truncated — 4649 more characters]
```

[4 more indexed source files omitted to keep this export small. The full file list is in the Codebase structure section above.]