Update app.py

app.py

@@ -57,6 +57,25 @@ pipe = DiffusionPipeline.from_pretrained(
 MAX_SEED = np.iinfo(np.int32).max
 
 
+def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+
+    image = pipe(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        guidance_scale=0.0
+    ).images[0]
+
+    return image, seed
+
+
+
+
 class MultiAgentState(BaseModel):
     state: List[str] = []
     messages: Annotated[list[AnyMessage], add_messages]