minor

app_v3.py

@@ -1,7 +1,7 @@
-import torch
 import subprocess
 subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
 
+import torch
 import spaces
 import os
 import moondream as md
@@ -46,45 +46,45 @@ 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)
+# @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")
+#     output = cap_model.generate(**generate_kwargs)
+#     print(f"Generated {len(output[0])} tokens")
 
 @spaces.GPU(duration=10)
 @torch.no_grad()
@@ -226,10 +226,10 @@ with gr.Blocks(title="FLUX Turbo Upscaler", fill_height=True) as iface:
         outputs=[output_caption, generated_image]
     )
     
-    caption_button.click(
-        fn=caption,
-        inputs=[control_image, system_prompt, temperature_slider, top_p_slider, max_tokens_slider, log_prompt],
-        outputs=output_caption,
-    )
+    # caption_button.click(
+    #     fn=caption,
+    #     inputs=[control_image, system_prompt, temperature_slider, top_p_slider, max_tokens_slider, log_prompt],
+    #     outputs=output_caption,
+    # )
 
-iface.launch(share=True)
+iface.launch()