incorporate InvokeAI into start_pipeline.py

3d-generation-pipeline/generate_image_local.py

@@ -1,28 +1,81 @@
-import torch
-from diffusers import StableDiffusionPipeline, StableDiffusion3Pipeline
-import time
+import requests
 
-start_timestamp = time.time()
-#model = "stabilityai/stable-diffusion-3.5-medium" # generation time: 13 min
-model = "stabilityai/stable-diffusion-3-medium-diffusers" # generation time: 10 min
-#model = "stabilityai/stable-diffusion-2" # generation time: 4 sec
+from invokeai_mcp_server import create_text2img_graph, enqueue_graph, wait_for_completion, get_image_url
+from urllib.parse import urljoin
 
-pipe = StableDiffusion3Pipeline.from_pretrained(model, torch_dtype=torch.float16)
-#pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
-pipe = pipe.to("cuda")
 
-model_loaded_timestamp = time.time()
-model_load_time = model_loaded_timestamp - start_timestamp
-print(f"model load time: {round(model_load_time)} seconds")
+INVOKEAI_BASE_URL = "http://127.0.0.1:9090"
 
-prompt = "A majestic broadsword with a golden pommel, no background"
-image = pipe(
-    prompt,
-    guidance_scale=3.0,
-).images[0]
 
-image_name = "image7.png"
-image.save(f"images/{image_name}")
+async def generate_image(arguments: dict):
 
-generation_time = time.time() - model_loaded_timestamp
-print(f"image generation time: {round(generation_time)} seconds")
+    # Extract parameters
+    prompt = arguments["prompt"]
+    negative_prompt = arguments.get("negative_prompt", "")
+    width = arguments.get("width", 512)
+    height = arguments.get("height", 512)
+    steps = arguments.get("steps", 30)
+    cfg_scale = arguments.get("cfg_scale", 7.5)
+    scheduler = arguments.get("scheduler", "euler")
+    seed = arguments.get("seed")
+    model_key = arguments.get("model_key")
+    lora_key = arguments.get("lora_key")
+    lora_weight = arguments.get("lora_weight", 1.0)
+    vae_key = arguments.get("vae_key")
+
+    print(f"Generating image with prompt: {prompt[:50]}...")
+
+    # Create graph
+    graph = await create_text2img_graph(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        model_key=model_key,
+        lora_key=lora_key,
+        lora_weight=lora_weight,
+        vae_key=vae_key,
+        width=width,
+        height=height,
+        steps=steps,
+        cfg_scale=cfg_scale,
+        scheduler=scheduler,
+        seed=seed
+    )
+
+    # Enqueue and wait for completion
+    result = await enqueue_graph(graph)
+    batch_id = result["batch"]["batch_id"]
+
+    print(f"Enqueued batch {batch_id}, waiting for completion...")
+
+    completed = await wait_for_completion(batch_id)
+
+    # Extract image name from result
+    if "result" in completed and "outputs" in completed["result"]:
+        outputs = completed["result"]["outputs"]
+        # Find the image output
+        for node_id, output in outputs.items():
+            if output.get("type") == "image_output":
+                image_name = output["image"]["image_name"]
+                image_url = await get_image_url(image_name)
+
+                return urljoin(INVOKEAI_BASE_URL, image_url)
+    
+    raise RuntimeError("Failed to generate image!")
+
+def download_file(url, filepath):
+    response = requests.get(url)
+
+    if response.status_code == 200:
+        with open(filepath, "wb") as file:
+            file.write(response.content)
+    else:
+        raise RuntimeError(f"Failed to download image. Status code: {response.status_code}")
+
+
+async def text_to_image_invoke_ai(prompt, output_path):
+    args = {
+        "prompt": prompt
+    }
+    image_url = await generate_image(args)
+    print("got image url: ", image_url)
+    download_file(image_url, output_path)