Prompt upsampling for diffusion models

import dspy

upsampler_llm = dspy.OpenAI(
    model="gpt-4",
    system_prompt="""
You are part of a team of bots that creates images. You work with an assistant bot that will draw anything
you say in square brackets. For example, outputting "a beautiful morning in the woods with the sun peaking
through the trees" will trigger your partner bot to output an image of a forest morning, as described.
You will be prompted by people looking to create detailed, amazing images. The way to accomplish this is to
take their short prompts and make them extremely detailed and descriptive.

There are a few rules to follow:
- You will only ever output a single image description per user request.
- Often times, the base prompt might consist of spelling mistakes or grammatical errors. You should correct
    such errors before making them extremely detailed and descriptive.
- Image descriptions must be between 15-80 words. Extra words will be ignored.
""",
)

class PromptUpsamplingSignature(dspy.Signature):
    base_prompt = dspy.InputField()
    answer = dspy.OutputField(
        desc="Create an imaginative image descriptive caption for the given base prompt."
    )

reasoning_attemps = [
    dspy.Prediction(
        rationale="a man holding a sword",
        answer="a pale figure with long white hair stands in the center of a dark forest, holding a sword high above his head.",
    ),
    dspy.Prediction(
        rationale="a frog playing dominoes",
        answer="a frog sits on a worn table playing a game of dominoes with an elderly raccoon. the table is covered in a green cloth, and the frog is wearing a jacket and a pair of jeans. The scene is set in a forest, with a large tree in the background.",
    ),
    dspy.Prediction(
        rationale="A bird scaring a scarecrow",
        answer="A large, vibrant bird with an impressive wingspan swoops down from the sky, letting out a piercing call as it approaches a weathered scarecrow in a sunlit field. The scarecrow, dressed in tattered clothing and a straw hat, appears to tremble, almost as if it's coming to life in fear of the approaching bird.",
    ),
    # ... you can implement more reasoning attempts for better result
]

prompt_upsampling_module = dspy.MultiChainComparison(
    PromptUpsamplingSignature, M=len(reasoning_attemps)
)

import torch
from diffusers import AutoPipelineForText2Image, DiffusionPipeline


class PromptUpsamplingModel(weave.Model):

    @weave.op()
    def predict(self, base_prompt) -> dict:
        with dspy.context(lm=upsampler_llm):
            return prompt_upsampling_module(
                reasoning_attemps, base_prompt=base_prompt
            ).answer


class StableDiffusionXLModel(weave.Model):
    diffusion_model: str
    enable_cpu_offload: bool = True
    prompt_upsampler: PromptUpsamplingModel
    _pipeline: DiffusionPipeline

    def __init__(
        self,
        diffusion_model: str,
        enable_cpu_offload: bool,
        prompt_upsampler: PromptUpsamplingModel
    ):
        super().__init__(
            diffusion_model=diffusion_model,
            enable_cpu_offload=enable_cpu_offload,
            prompt_upsampler=prompt_upsampler,
        )
        self._pipeline = AutoPipelineForText2Image.from_pretrained(
            self.diffusion_model,
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True,
        )
        if self.enable_cpu_offload:
            self._pipeline.enable_model_cpu_offload()
        else:
            self._pipeline = self._pipeline.to("cuda")
    
    @weave.op()
    def predict(
        self,
        base_prompt: str,
        negative_prompt: Optional[str] = None,
        num_inference_steps: Optional[int] = 50,
        image_size: Optional[int] = 1024,
        guidance_scale: Optional[float] = 7.0,
    ) -> dict:
        upsampled_prompt = self.prompt_upsampler.predict(base_prompt)
        image = self._pipeline(
            prompt=upsampled_prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            height=image_size,
            width=image_size,
            guidance_scale=guidance_scale,
        ).images[0]
        return {
            "upsampled_prompt": upsampled_prompt,
            "image": base64_encode_image(image)
        }


prompt_upsampler=PromptUpsamplingModel()

model = StableDiffusionXLModel(
    diffusion_model="stabilityai/stable-diffusion-xl-base-1.0",
    enable_cpu_offload=True,
    prompt_upsampler=prompt_upsampler,
)

sdxl_prediction = model.predict(base_prompt="a frog dressed as a knight")

from dsp import GPT3
from openai import OpenAI


class DSPyOpenAIMultiModalLM(GPT3):

    def __init__(
        self,
        model: str = "gpt-4o",
        api_key: str | None = None,
        system_prompt: str | None = None,
        **kwargs,
    ):
        super().__init__(
            model,
            api_key,
            api_provider="openai",
            api_base=None,
            model_type=None,
            system_prompt=system_prompt,
            **kwargs,
        )
        self.model_type = model
        self._openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

    @weave.op()
    def create_messages(self, prompt: str):
        images = find_base64_images(prompt)
        for image in images:
            prompt = prompt.replace(image, "")

        user_prompt = [{"type": "text", "text": prompt}]
        for image in images:
            user_prompt.append({"type": "image_url", "image_url": {"url": image}})
        messages = []
        if self.system_prompt:
            messages.append({"role": "system", "content": self.system_prompt})
        messages.append({"role": "user", "content": user_prompt})
        return messages

    @weave.op()
    def basic_request(self, prompt: str, **kwargs):
        messages = self.create_messages(prompt)
        response = self._openai_client.chat.completions.create(
            model=self.model_type, messages=messages, **kwargs
        )
        self.history.append({"prompt": prompt, "response": response, "kwargs": kwargs})
        return response

    @weave.op()
    def request(self, prompt: str, **kwargs):
        return super().request(prompt, **kwargs)

    @weave.op()
    def __call__(
        self, prompt: str, only_completed: bool = True, **kwargs
    ) -> list:
        response = self.request(prompt, **kwargs)
        choices = (
            [choice for choice in response.choices if choice.finish_reason == "stop"]
            if only_completed and len(response.choices) != 0
            else response.choices
        )
        return [choice.message.content for choice in choices]

class JudgeInput(BaseModel):
    base_prompt: str = Field(description="The base prompt used to generate the image")
    generated_image: str = Field(description="The generated image")


class JudgeMent(BaseModel):
    think_out_loud: str = Field(
        description="Think out loud about your eventual judgement"
    )
    score: float = Field(description="A score between 0 and 1")
    judgement: str = Field(description="Output either 'correct' or 'incorrect'")


class JudgeSignature(dspy.Signature):
    input: JudgeInput = dspy.InputField()
    output: JudgeMent = dspy.OutputField()


class MultiModalJudgeModule(dspy.Module):

    def __init__(self):
        self.prog = dspy.TypedPredictor(JudgeSignature)

    @weave.op()
    def forward(self, base_prompt: str, generated_image: str) -> dict:
        return self.prog(
            input=JudgeInput(
                base_prompt=base_prompt, generated_image=generated_image
            )
        ).output


judgement_module = MultiModalJudgeModule()

JUDGE_SYSTEM_PROMPT = """
You are responsible for judging the faithfulness of images generated by a computer program to the
base prompt used to generate them. You will be presented with an image and given the base prompt
that was used to produce the image. The base prompts you are judging are designed to stress-test
image generation programs, and may include things such as:
1. Scrambled or mis-spelled words (the image generator should an image associated with
    the probably meaning).
2. Color assignment (the image generator should apply the correct color to the correct object).
3. Color assignment (the image generator should apply the correct color to the correct object).
4. Abnormal associations, for example 'elephant under a sea', where the image should depict
    what is requested.
5. Descriptions of objects, the image generator should draw the most commonly associated object.
6. Rare single words, where the image generator should create an image somewhat associable with
    the specified image.
7. Images with text in them, where the image generator should create an image with the specified
    text in it. You need to make a decision as to whether or not the image is correct, given the
    base prompt.

You will first think out loud about your eventual judgement, enumerating reasons why the image
does or does not match the given base prompt. After thinking out loud, you should assign a score
between 0 and 1 depending on how much you think the image is faithful to the base prompt. Next,
you should output either 'correct' or 'incorrect' depending on whether you think the image is
faithful to the base prompt.

A few rules:
1. The score should be used to indicate how close the image is to the base prompt in terms of objects,
    color or count; with 0 being very far and 1 being very close.
2. If other objects are present in the image that are not explicitly mentioned by the base prompt,
    assign a higher score.
3. If the objects being displayed is deformed, assign a lower score. Assign a higher score, if the objects
    are displayed in a more detailed manner.
4. 'incorrect' should be reserved for instances where a specific aspect of the base prompt is not followed
    correctly, such as a wrong object, color or count and the score should be less than or equal to 0.5.
"""

class OpenAIJudgeModel(weave.Model):
    openai_model: str
    seed: int
    _judgement_llm: dspy.Module

    def __init__(self, openai_model: str = "gpt-4-turbo", seed: int = 42):
        super().__init__(openai_model=openai_model, seed=seed)
        self._judgement_llm = DSPyOpenAIMultiModalLM(
            model="gpt-4o", system_prompt=JUDGE_SYSTEM_PROMPT, seed=self.seed
        )

    @weave.op()
    def predict(self, base_prompt: str, generated_image: str) -> JudgeMent:
        with dspy.context(lm=self._judgement_llm):
            judgement = judgement_module(base_prompt, generated_image)
        return judgement

    @weave.op()
    def score(self, base_prompt: str, model_output: dict) -> dict:
        judgement: JudgeMent = self.predict(
            base_prompt=base_prompt, generated_image=model_output["image"]
        )
        return {
            "score": judgement.score,
            "is_image_correct": judgement.judgement == "correct",
        }

judge_model = OpenAIJudgeModel()
judgement = judge_model.score("a frog dressed as a knight", sdxl_prediction)