test: focus area

app_v3.py

@@ -21,12 +21,6 @@ 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)
 
 md_api_key = os.getenv("MD_KEY")
 model = md.vl(api_key=md_api_key)
@@ -46,46 +40,6 @@ pipe = FluxControlNetPipeline.from_pretrained(
 )
 pipe.to("cuda")
 
-# @spaces.GPU()
-# @torch.no_grad()
-# def caption(input_image: Image.Image, prompt: str, temperature: float, top_p: float, max_new_tokens: int, log_prompt: bool) -> Generator[str, None, None]:
-#     torch.cuda.empty_cache()
-#     if input_image is None:
-#         yield "No image provided. Please upload an image."
-#         return
-#     if log_prompt:
-#         print(f"PromptLog: {repr(prompt)}")
-#     convo = [
-#         {
-#             "role": "system",
-#             "content": "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions.",
-#         },
-#         {
-#             "role": "user",
-#             "content": prompt.strip(),
-#         },
-#     ]
-#     convo_string = cap_processor.apply_chat_template(convo, tokenize=False, add_generation_prompt=True)
-#     assert isinstance(convo_string, str)
-#     inputs = cap_processor(text=[convo_string], images=[input_image], return_tensors="pt").to('cuda')
-#     inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
-#     streamer = TextIteratorStreamer(cap_processor.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-#     generate_kwargs = dict(
-#         **inputs,
-#         max_new_tokens=max_new_tokens,
-#         do_sample=True if temperature > 0 else False,
-#         suppress_tokens=None,
-#         use_cache=True,
-#         temperature=temperature if temperature > 0 else None,
-#         top_k=None,
-#         top_p=top_p if temperature > 0 else None,
-#         streamer=streamer,
-#     )
-#     _= cap_model.generate(**generate_kwargs)
-
-#     output = cap_model.generate(**generate_kwargs)
-#     print(f"Generated {len(output[0])} tokens")
-
 @spaces.GPU(duration=10)
 @torch.no_grad()
 def generate_image(prompt, scale, steps, control_image, controlnet_conditioning_scale, guidance_scale, seed, guidance_end):
@@ -124,6 +78,17 @@ def generate_caption(control_image):
     
     return detailed_caption
 
+def generate_focus(control_image, focus_list):
+    if control_image is None:
+        return None, None
+    
+    # Generate a detailed caption
+    focus_query = model.query(control_image, "Please provide a concise but illustrative description of the following area(s) of focus: " + focus_list)
+    focus_description = focus_query["answer"]
+    print(f"Areas of focus: {focus_description}") 
+
+    return focus_description
+
 def process_image(control_image, user_prompt, system_prompt, scale, steps, 
                 controlnet_conditioning_scale, guidance_scale, seed, 
                 guidance_end, temperature, top_p, max_new_tokens, log_prompt):
@@ -133,7 +98,7 @@ 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")
+        mcaption = model.caption(control_image, length="normal")
         detailed_caption = mcaption["caption"]
         print(f"Detailed caption: {detailed_caption}")
 
@@ -172,11 +137,12 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
         generated_image = gr.Image(type="pil", label="Generated Image", format="png", show_label=False)
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Textbox(lines=4, placeholder="Enter your prompt here...", label="Prompt")
-            output_caption = gr.Textbox(label="Caption")
-            scale = gr.Slider(1, 3, value=1, label="Scale", step=0.25)
-            generate_button = gr.Button("Generate Image", variant="primary")
-            caption_button = gr.Button("Generate Caption", variant="secondary")
+            prompt = gr.Textbox(lines=4, info="Enter your prompt here or wait for auto-generation...", label="Image Description")
+            focus = gr.Textbox(label="Areas of Focus", info="e.g. 'face', 'eyes', 'hair', 'clothes', 'background', etc.")
+            scale = gr.Slider(1, 3, value=1, label="Scale (Upscale Factor)", step=0.25)
+            with gr.Row():
+                generate_button = gr.Button("Generate Image", variant="primary")
+                caption_button = gr.Button("Generate Caption", variant="secondary")
         with gr.Column(scale=1):
             seed = gr.Slider(0, MAX_SEED, value=42, label="Seed", step=1)
             steps = gr.Slider(2, 16, value=8, label="Steps", step=1)
@@ -219,12 +185,17 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as demo:
             controlnet_conditioning_scale, guidance_scale, seed, 
             guidance_end, temperature_slider, top_p_slider, max_tokens_slider, log_prompt
         ],
-        outputs=[output_caption, generated_image, prompt]
+        outputs=[generated_image, prompt]
     )
     control_image.change(
         generate_caption,
         inputs=[control_image],
         outputs=[prompt]
+    ).then(
+        generate_focus,
+        inputs=[control_image, focus],
+        outputs=[prompt],
+        _js="(caption, focus) => caption + ' ' + focus"
     )
     caption_button.click(
         fn=generate_caption,