Update ominicontrol.py

ominicontrol.py

@@ -3,15 +3,14 @@ from diffusers.pipelines import FluxPipeline
 from OminiControl.src.flux.condition import Condition
 from PIL import Image
 import random
-import os
 
 from OminiControl.src.flux.generate import generate, seed_everything
 
-HF_TOKEN=os.getenv("HF_TOKEN")
+from log import insert_log, log_image
 
 print("Loading model...")
 pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, use_auth_token=HF_TOKEN
+    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16
 )
 pipe = pipe.to("cuda")
 
@@ -21,29 +20,24 @@ pipe.load_lora_weights(
     "Yuanshi/OminiControlStyle",
     weight_name=f"v0/ghibli.safetensors",
     adapter_name="ghibli",
-    token=HF_TOKEN
 )
 pipe.load_lora_weights(
     "Yuanshi/OminiControlStyle",
     weight_name=f"v0/irasutoya.safetensors",
     adapter_name="irasutoya",
-    token=HF_TOKEN
 )
 pipe.load_lora_weights(
     "Yuanshi/OminiControlStyle",
     weight_name=f"v0/simpsons.safetensors",
     adapter_name="simpsons",
-    token=HF_TOKEN
 )
 pipe.load_lora_weights(
     "Yuanshi/OminiControlStyle",
     weight_name=f"v0/snoopy.safetensors",
     adapter_name="snoopy",
-    token=HF_TOKEN
 )
 
 
-
 def generate_image(
     image,
     style,
@@ -69,7 +63,7 @@ def generate_image(
         "The Simpsons": "simpsons",
         "Snoopy": "snoopy",
     }[style]
-    pipe.set_adapters(activate_adapter_name)
+    # pipe.set_adapters(activate_adapter_name)
 
     factor = 512 / max(image.size)
     image = resize(
@@ -125,5 +119,39 @@ def generate_image(
         default_lora=True,
         max_sequence_length=32,
     ).images[0]
+    # result_img = image
+
+    condition_id = log_image(result_img)
+    result_id = log_image(result_img)
+
+    log_data = {
+        "condition": condition_id,
+        "result": result_id,
+        "prompt": "",
+        "inference_mode": inference_mode,
+        "image_guidance_scale": image_guidance,
+        "seed": seed,
+        "steps": steps,
+        "width": width,
+        "height": height,
+    }
+    log_data = {k: str(v) for k, v in log_data.items()}
+
+    _, log_id = insert_log("inference", log_data)
+
+    print(f"Image log ID: {log_id}")
+
+    return result_img, log_id
+
+
+def vote_feedback(
+    log_id,
+    feedback,
+):
+    log_data = {
+        "log_id": log_id,
+        "feedback": feedback,
+    }
+    log_data = {k: str(v) for k, v in log_data.items()}
 
-    return result_img
+    insert_log("feedback", log_data)