dy2000/optimized-diffusers-code

1import google.generativeai as genai2import os3 4# https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview5MODEL_NAME = "gemini-2.5-flash-preview-05-20"6 7 8class LLMCodeOptimizer:9    def __init__(self, model_name: str, system_prompt: str):10        api_key = os.getenv("GOOGLE_API_KEY")11        if api_key is None:12            raise ValueError("Must provide an API key for Gemini through the `GOOGLE_API_KEY` env variable.")13        genai.configure(api_key=api_key)14        self.model = genai.GenerativeModel(model_name=model_name, system_instruction=system_prompt)15 16    def __call__(self, generation_prompt):17        try:18            print("Sending request to Gemini...")19            response = self.model.generate_content(generation_prompt)20 21            return response.text22 23        except Exception as e:24            # Handle potential exceptions, such as invalid API keys,25            # network issues, or content moderation errors.26            return f"An error occurred: {e}"27