Update app.py

app.py

@@ -39,6 +39,24 @@ import re, operator
 
 chromadb.api.client.SharedSystemClient.clear_system_cache()
 
+
+import torch
+from diffusers import DiffusionPipeline
+import numpy as np
+import random
+
+# Load FLUX.1-schnell model
+torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch_dtype
+).to(device)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+
 class MultiAgentState(BaseModel):
     state: List[str] = []
     messages: Annotated[list[AnyMessage], add_messages]