test: new cap

app_v3.py

@@ -1,6 +1,10 @@
 import torch
+import subprocess
+subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+
 import spaces
 import os
+import moondream as md
 from diffusers.utils import load_image
 from diffusers.hooks import apply_group_offloading
 from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
@@ -17,12 +21,15 @@ import gradio as gr
 
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 MAX_SEED = 1000000
-MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
-cap_processor = AutoProcessor.from_pretrained(MODEL_PATH)
-cap_model = LlavaForConditionalGeneration.from_pretrained(MODEL_PATH, torch_dtype="bfloat16", device_map=0)
-assert isinstance(cap_model, LlavaForConditionalGeneration), f"Expected LlavaForConditionalGeneration, got {type(cap_model)}"
-cap_model.eval()
-apply_liger_kernel_to_llama(model=cap_model.language_model)
+# MODEL_PATH = "fancyfeast/llama-joycaption-beta-one-hf-llava"
+# cap_processor = AutoProcessor.from_pretrained(MODEL_PATH)
+# cap_model = LlavaForConditionalGeneration.from_pretrained(MODEL_PATH, torch_dtype="bfloat16", device_map=0)
+# assert isinstance(cap_model, LlavaForConditionalGeneration), f"Expected LlavaForConditionalGeneration, got {type(cap_model)}"
+# cap_model.eval()
+# apply_liger_kernel_to_llama(model=cap_model.language_model)
+
+md_api_key = os.getenv("MD_KEY")
+model = md.vl(api_key=md_api_key)
 
 text_encoder_2_unquant = T5EncoderModel.from_pretrained(
     "LPX55/FLUX.1-merged_uncensored",
@@ -114,23 +121,27 @@ def process_image(control_image, user_prompt, system_prompt, scale, steps,
     
     # If no user prompt provided, generate a caption first
     if not final_prompt:
+        # Generate a detailed caption
+        mcaption = model.caption(control_image, length="long")
+        detailed_caption = mcaption["caption"]
+        print(f"Detailed caption: {detailed_caption}")
         # Generate caption
-        caption_gen = caption(
-            input_image=control_image,
-            prompt=system_prompt,
-            temperature=temperature,
-            top_p=top_p,
-            max_new_tokens=max_new_tokens,
-            log_prompt=log_prompt
-        )
+        # caption_gen = caption(
+        #     input_image=control_image,
+        #     prompt=system_prompt,
+        #     temperature=temperature,
+        #     top_p=top_p,
+        #     max_new_tokens=max_new_tokens,
+        #     log_prompt=log_prompt
+        # )
         
         # Get the full caption by exhausting the generator
-        generated_caption = ""
-        for chunk in caption_gen:
-            generated_caption += chunk
-            yield generated_caption, None  # Update caption in real-time
+        # generated_caption = ""
+        # for chunk in caption_gen:
+        #     generated_caption += chunk
+        #     yield generated_caption, None  # Update caption in real-time
         
-        final_prompt = generated_caption
+        final_prompt = detailed_caption
         yield f"Using caption: {final_prompt}", None
     
     # Show the final prompt being used