Update app.py

app.py

@@ -1,64 +1,29 @@
-import os
-import time
-import numpy as np
-from PIL import Image
-
-# --- Hugging Face and Diffusers Imports ---
 import huggingface_hub
 # Monkey-patch: if cached_download is missing, alias it to hf_hub_download.
 if not hasattr(huggingface_hub, "cached_download"):
     huggingface_hub.cached_download = huggingface_hub.hf_hub_download
+
 print("huggingface_hub version:", huggingface_hub.__version__)
 
 import diffusers
 print("diffusers version:", diffusers.__version__)
 import numpy
 print("numpy version:", numpy.__version__)
+import spaces
 import gradio as gr
 import torch
 from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline, AutoencoderKL
+from PIL import Image
+import os
+import time
+
+from utils.dl_utils import dl_cn_model, dl_cn_config, dl_tagger_model, dl_lora_model
+from utils.image_utils import resize_image_aspect_ratio, base_generation
 
-# --- Additional Utilities and Download Functions ---
-# (Assuming these are defined elsewhere in your project. For integration, dummy definitions are provided.)
-def dl_cn_model(cn_dir):
-    # Dummy: in your project, this downloads the ControlNet model
-    print("Downloading ControlNet model to", cn_dir)
-def dl_cn_config(cn_dir):
-    print("Downloading ControlNet config to", cn_dir)
-def dl_tagger_model(tagger_dir):
-    print("Downloading tagger model to", tagger_dir)
-def dl_lora_model(lora_dir):
-    print("Downloading LoRA model to", lora_dir)
-    
-# Dummy image utility functions
-def resize_image_aspect_ratio(image):
-    # For demonstration, return the image as is.
-    return image
-
-def base_generation(size, color):
-    # Create a blank image with the given color and size.
-    return Image.new("RGBA", size, color)
-
-# Dummy prompt utilities
-def execute_prompt(tags, prompt):
-    # In your project, this may combine tags and prompt
-    return prompt + " " + ", ".join(tags)
-def remove_color(prompt):
-    # Dummy function: simply return prompt unchanged.
-    return prompt
-def remove_duplicates(prompt):
-    # Dummy function: simply return prompt unchanged.
-    return prompt
-
-# Dummy tagger function (if needed)
-def modelLoad(tagger_dir):
-    # Return a dummy model
-    return None
-def analysis(image_path, tagger_dir, tagger_model):
-    # Return dummy tags
-    return ["lineart", "sketch"]
-
-# --- Set Up Directories ---
+from utils.prompt_utils import execute_prompt, remove_color, remove_duplicates
+from utils.tagger import modelLoad, analysis
+
+# Set up directories
 path = os.getcwd()
 cn_dir = os.path.join(path, "controlnet")
 tagger_dir = os.path.join(path, "tagger")
@@ -67,43 +32,46 @@ os.makedirs(cn_dir, exist_ok=True)
 os.makedirs(tagger_dir, exist_ok=True)
 os.makedirs(lora_dir, exist_ok=True)
 
-# --- Download Required Models/Configs ---
+# Download required models and configs
 dl_cn_model(cn_dir)
 dl_cn_config(cn_dir)
 dl_tagger_model(tagger_dir)
 dl_lora_model(lora_dir)
 
-# --- Diffusers-Based Model Loading and Predict Function ---
 def load_model(lora_dir, cn_dir):
     dtype = torch.float16
-    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype)
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
     controlnet = ControlNetModel.from_pretrained(cn_dir, torch_dtype=dtype, use_safetensors=True)
+
     pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
-        "cagliostrolab/animagine-xl-3.1", controlnet=controlnet, vae=vae, torch_dtype=dtype
+        "cagliostrolab/animagine-xl-3.1", controlnet=controlnet, vae=vae, torch_dtype=torch.float16
     )
     pipe.enable_model_cpu_offload()
-    # Load your LoRA weights (assumes the file "lineart.safetensors" is in lora_dir)
     pipe.load_lora_weights(lora_dir, weight_name="lineart.safetensors")
     return pipe
 
+@spaces.GPU(duration=120)
 def predict(input_image_path, prompt, negative_prompt, controlnet_scale):
-    pipe = load_model(lora_dir, cn_dir)
-    input_image = Image.open(input_image_path).convert("RGB")
+    pipe = load_model(lora_dir, cn_dir) 
+    input_image = Image.open(input_image_path)
     base_image = base_generation(input_image.size, (255, 255, 255, 255)).convert("RGB")
-    resize_img = resize_image_aspect_ratio(input_image)
-    resize_base = resize_image_aspect_ratio(base_image)
+    resize_image = resize_image_aspect_ratio(input_image)
+    resize_base_image = resize_image_aspect_ratio(base_image)
     generator = torch.manual_seed(0)
     last_time = time.time()
-    
-    # Here we assume the prompt is already generated (or modified) by the new Janus logic.
-    final_prompt = prompt  # Optionally, prepend additional base phrases if needed.
-    print("Final prompt:", final_prompt)
+    # Prepend a base prompt to get best results
+    prompt = "masterpiece, best quality, monochrome, sharp uniform black lines, vector style, very thick lineart, clean lineart, no shading, solid very thick black lines, no gradients, white background, " + prompt
+    execute_tags = ["sketch", "transparent background"]
+    prompt = execute_prompt(execute_tags, prompt)
+    prompt = remove_duplicates(prompt)        
+    prompt = remove_color(prompt)
+    print(prompt)
 
     output_image = pipe(
-        image=resize_base,
-        control_image=resize_img,
+        image=resize_base_image,
+        control_image=resize_image,
         strength=1.0,
-        prompt=final_prompt,
+        prompt=prompt,
         negative_prompt=negative_prompt,
         controlnet_conditioning_scale=float(controlnet_scale),
         generator=generator,
@@ -114,91 +82,35 @@ def predict(input_image_path, prompt, negative_prompt, controlnet_scale):
     output_image = output_image.resize(input_image.size, Image.LANCZOS)
     return output_image
 
-# --- Janus-Based Prompt Generation Function ---
-# This code is taken from your second Hugging Face app.
-from transformers import AutoConfig, AutoModelForCausalLM
-from janus.models import VLChatProcessor
-
-# Set up Janus model and processor
-janus_model_path = "deepseek-ai/Janus-Pro-1B"
-janus_config = AutoConfig.from_pretrained(janus_model_path)
-language_config = janus_config.language_config
-language_config._attn_implementation = 'eager'
-vl_gpt = AutoModelForCausalLM.from_pretrained(
-    janus_model_path,
-    language_config=language_config,
-    trust_remote_code=True
-)
-if torch.cuda.is_available():
-    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
-else:
-    vl_gpt = vl_gpt.to(torch.float16)
-vl_chat_processor = VLChatProcessor.from_pretrained(janus_model_path)
-tokenizer = vl_chat_processor.tokenizer
-cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-@torch.inference_mode()
-@gr.analytics.track()  # Using gradio decorator for potential GPU allocation if using spaces; you can remove if not needed.
-def generate_prompt_from_image(input_image_path, seed: int = 42, top_p: float = 0.95, temperature: float = 0.1):
-    # Open image
-    image = Image.open(input_image_path).convert("RGB")
-    image_np = np.array(image)
-    
-    # Set seed for reproducibility
-    torch.cuda.empty_cache()
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-    
-    # Construct conversation for the Janus model
-    conversation = [
-        {
-            "role": "<|User|>",
-            "content": (
-                "<image_placeholder>\nGenerate a detailed artistic prompt for extracting crisp, high-quality lineart. "
-                "The prompt should include phrases like 'masterpiece, best quality, monochrome, sharp uniform black lines, "
-                "vector style, very thick lineart, clean lineart, no shading, solid very thick black lines, no gradients, white background'."
-            ),
-            "images": [image_np],
-        },
-        {"role": "<|Assistant|>", "content": ""},
-    ]
-    pil_images = [image]
-    prepare_inputs = vl_chat_processor(
-        conversations=conversation, images=pil_images, force_batchify=True
-    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
-    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
-    
-    outputs = vl_gpt.language_model.generate(
-        inputs_embeds=inputs_embeds,
-        attention_mask=prepare_inputs.attention_mask,
-        pad_token_id=tokenizer.eos_token_id,
-        bos_token_id=tokenizer.bos_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        max_new_tokens=64,
-        do_sample=True if temperature != 0 else False,
-        use_cache=True,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    
-    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
-    return answer
+@spaces.GPU(duration=120)
+def prompt_analysis(input_image_path):
+    """
+    Run prompt analysis on the given image.
+    Loads the tagger model, runs analysis, cleans the tags, and returns a string.
+    """
+    # Load the tagger model using the tagger_dir (set earlier in the file)
+    tagger_model = modelLoad(tagger_dir)
+    tags = analysis(input_image_path, tagger_dir, tagger_model)
+    tags_clean = remove_color(tags)
+    if isinstance(tags_clean, (list, tuple)):
+        return ", ".join(tags_clean)
+    return tags_clean
+
 
-# --- Gradio Interface ---
 class Img2Img:
     def __init__(self):
         self.demo = self.layout()
-    
-    def generate_prompt_callback(self, input_image_path):
-        # Generate prompt using Janus model
-        return generate_prompt_from_image(input_image_path)
-    
-    def generate_image_callback(self, input_image_path, prompt, negative_prompt, controlnet_scale):
-        # Use the provided prompt (which may have been auto-generated) to generate the image
-        return predict(input_image_path, prompt, negative_prompt, controlnet_scale)
-    
+        self.tagger_model = None
+        self.input_image_path = None
+        self.canny_image = None
+
+    def process_prompt_analysis(self, input_image_path):
+        if self.tagger_model is None:
+            self.tagger_model = modelLoad(tagger_dir)
+        tags = analysis(input_image_path, tagger_dir, self.tagger_model)
+        tags_list = remove_color(tags)
+        return tags_list
+
     def layout(self):
         css = """
         #intro{
@@ -210,34 +122,35 @@ class Img2Img:
         with gr.Blocks(css=css) as demo:
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(label="Input Image", type="filepath")
-                    prompt = gr.Textbox(label="Prompt", lines=3)
-                    negative_prompt = gr.Textbox(
-                        label="Negative Prompt",
-                        lines=3,
-                        value="sketch, lowres, error, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, blurry"
-                    )
-                    controlnet_scale = gr.Slider(
-                        minimum=0.5, maximum=1.25, value=1.0, step=0.01, label="Lineart Fidelity"
+                    self.input_image_path = gr.Image(label="Input image", type='filepath')
+                    self.prompt = gr.Textbox(label="Prompt", lines=3)
+                    self.negative_prompt = gr.Textbox(
+                        label="Negative prompt", 
+                        lines=3, 
+                        value="sketch, lowres, error, extra digit, fewer digits, cropped, worst quality,low quality, normal quality, jpeg artifacts, blurry"
                     )
-                    generate_prompt_btn = gr.Button("Generate Prompt")
-                    generate_image_btn = gr.Button("Generate Image", variant="primary")
+                    # Button to run prompt analysis locally (UI callback)
+                    prompt_analysis_button = gr.Button("Prompt analysis")
+                    self.controlnet_scale = gr.Slider(
+                        minimum=0.5, maximum=1.25, value=1.0, step=0.01, label="Lineart fidelity"
+                    )                 
+                    generate_button = gr.Button(value="Generate", variant="primary")
                 with gr.Column():
-                    output_image = gr.Image(type="pil", label="Output Image")
-            
-            # Button callbacks
-            generate_prompt_btn.click(
-                fn=self.generate_prompt_callback,
-                inputs=[input_image],
-                outputs=prompt
+                    self.output_image = gr.Image(type="pil", label="Output image")
+
+            prompt_analysis_button.click(
+                self.process_prompt_analysis,
+                inputs=[self.input_image_path],
+                outputs=self.prompt
             )
-            generate_image_btn.click(
-                fn=self.generate_image_callback,
-                inputs=[input_image, prompt, negative_prompt, controlnet_scale],
-                outputs=output_image
+
+            generate_button.click(
+                fn=predict,
+                inputs=[self.input_image_path, self.prompt, self.negative_prompt, self.controlnet_scale],
+                outputs=self.output_image
             )
         return demo
 
 img2img = Img2Img()
 img2img.demo.queue()
-img2img.demo.launch(share=True, show_error=True)
+img2img.demo.launch(share=True, show_error=True)