forked from llmsresearch/paperbanana
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdebug_pipeline_run.py
More file actions
42 lines (28 loc) · 1.12 KB
/
debug_pipeline_run.py
File metadata and controls
42 lines (28 loc) · 1.12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
import asyncio
from PIL import Image
from paperbanana.core.pipeline import PaperBananaPipeline
from paperbanana.core.types import DiagramType, GenerationInput
# -------- Fake VLM --------
class FakeVLM:
name = "fake-vlm"
model_name = "fake-model"
async def generate(self, *args, **kwargs):
return "fake response"
# -------- Fake Image Generator --------
class FakeImageGen:
async def generate(self, prompt=None, output_path=None, iteration=None, seed=None, **kwargs):
iteration = iteration or 1 # fallback to 1 if None
img = Image.new("RGB", (256, 256), color=(iteration * 40 % 256, 100, 150))
return img
async def main():
pipeline = PaperBananaPipeline(vlm_client=FakeVLM(), image_gen_fn=FakeImageGen())
inp = GenerationInput(
source_context="A neural network with input, hidden and output layers.",
communicative_intent="Explain feedforward architecture",
diagram_type=DiagramType.METHODOLOGY,
raw_data=None,
)
output = await pipeline.generate(inp)
print("TIMING FOUND:")
print(output.metadata["timing"])
asyncio.run(main())