Update app.py

app.py

@@ -1,74 +1,104 @@
 import gradio as gr
-from PIL import Image
+import torch
+from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
+from diffusers import DiffusionPipeline
 import requests
-from diffusers import StableDiffusionPipeline
-
-# Load models using diffusers
-general_model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-anime_model = StableDiffusionPipeline.from_pretrained("hakurei/waifu-diffusion")
-
-# Placeholder functions for the actual implementations
-def check_anime_image(image):
-    # Use SauceNAO or similar service to check if the image is anime
-    # and fetch similar images and tags
-    return False, [], []
-
-def describe_image_general(image):
-    # Use the general model to describe the image
-    description = general_model(image)
-    return description
-
-def describe_image_anime(image):
-    # Use the anime model to describe the image
-    description = anime_model(image)
-    return description
+from PIL import Image
+from io import BytesIO
 
-def merge_tags(tags1, tags2):
-    # Merge tags, removing duplicates
-    return list(set(tags1 + tags2))
+# Initialize models
+anime_model = DiffusionPipeline.from_pretrained("SmilingWolf/wd-v1-4-vit-tagger")
+photo_model = AutoModelForZeroShotImageClassification.from_pretrained("facebook/florence-base-in21k-retrieval")
+processor = AutoProcessor.from_pretrained("facebook/florence-base-in21k-retrieval")
 
-# Gradio app functions
-def process_image(image, mode):
-    # Convert the image to a format suitable for the models
-    image = image.resize((256, 256))
+def get_booru_image(booru, image_id):
+    # This is a placeholder function. You'd need to implement the actual API calls for each booru.
+    url = f"https://api.{booru}.org/images/{image_id}"
+    response = requests.get(url)
+    img = Image.open(BytesIO(response.content))
+    tags = ["tag1", "tag2", "tag3"]  # Placeholder
+    return img, tags
 
-    if mode == "Anime":
-        is_anime, similar_images, original_tags = check_anime_image(image)
-        if is_anime:
-            tags = describe_image_anime(image)
-            return tags, original_tags
-        else:
-            return ["Not an anime image"], []
+def transcribe_image(image, image_type, transcriber, booru_tags=None):
+    if image_type == "Anime":
+        with torch.no_grad():
+            tags = anime_model(image)
     else:
-        tags = describe_image_general(image)
-        return tags, []
+        inputs = processor(images=image, return_tensors="pt")
+        outputs = photo_model(**inputs)
+        tags = outputs.logits.topk(50).indices.squeeze().tolist()
+        tags = [processor.config.id2label[t] for t in tags]
+    
+    if booru_tags:
+        tags = list(set(tags + booru_tags))
+    
+    return ", ".join(tags)
 
-def describe(image, mode):
-    tags, original_tags = process_image(image, mode)
-    return gr.update(value="\n".join(tags)), gr.update(value="\n".join(original_tags))
+def update_image(image_type, booru, image_id, uploaded_image):
+    if image_type == "Anime" and booru != "Upload":
+        image, booru_tags = get_booru_image(booru, image_id)
+        return image, gr.update(visible=True), booru_tags
+    elif uploaded_image is not None:
+        return uploaded_image, gr.update(visible=True), None
+    else:
+        return None, gr.update(visible=False), None
 
-def merge(tags, original_tags):
-    merged_tags = merge_tags(tags.split("\n"), original_tags.split("\n"))
-    return "\n".join(merged_tags)
+def on_image_type_change(image_type):
+    if image_type == "Anime":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(choices=["Anime", "Photo/Other"])
+    else:
+        return gr.update(visible=False), gr.update(visible=True), gr.update(choices=["Photo/Other", "Anime"])
 
-# Gradio interface
-with gr.Blocks() as demo:
-    with gr.Row():
-        image_input = gr.Image(type="pil", tool="editor", label="Upload/Paste Image")
-        mode = gr.Dropdown(choices=["Anime", "General"], label="Mode")
+with gr.Blocks() as app:
+    gr.Markdown("# Image Transcription App")
+    
+    with gr.Tab("Step 1: Image"):
+        image_type = gr.Dropdown(["Anime", "Photo/Other"], label="Image type")
+        
+        with gr.Column(visible=False) as anime_options:
+            booru = gr.Dropdown(["Gelbooru", "Danbooru", "Upload"], label="Boorus")
+            image_id = gr.Textbox(label="Image ID")
+            get_image_btn = gr.Button("Get image")
+        
+        upload_btn = gr.UploadButton("Upload Image", visible=False)
+        
+        image_display = gr.Image(label="Image to transcribe", visible=False)
+        booru_tags = gr.State(None)
+        
+        transcribe_btn = gr.Button("Transcribe", visible=False)
+        transcribe_with_tags_btn = gr.Button("Transcribe with booru tags", visible=False)
+    
+    with gr.Tab("Step 2: Transcribe"):
+        transcriber = gr.Dropdown(["Anime", "Photo/Other"], label="Transcriber")
+        transcribe_image_display = gr.Image(label="Image to transcribe")
+        transcribe_btn_final = gr.Button("Transcribe")
+        tags_output = gr.Textbox(label="Transcribed tags")
+    
+    image_type.change(on_image_type_change, inputs=[image_type], 
+                      outputs=[anime_options, upload_btn, transcriber])
+    
+    get_image_btn.click(update_image, 
+                        inputs=[image_type, booru, image_id, upload_btn], 
+                        outputs=[image_display, transcribe_btn, booru_tags])
+    
+    upload_btn.upload(update_image, 
+                      inputs=[image_type, booru, image_id, upload_btn], 
+                      outputs=[image_display, transcribe_btn, booru_tags])
     
-    describe_button = gr.Button("Describe")
-    merge_button = gr.Button("Merge Tags")
+    def transcribe_and_update(image, image_type, transcriber, booru_tags):
+        tags = transcribe_image(image, image_type, transcriber, booru_tags)
+        return image, tags
     
-    with gr.TabGroup() as tab_group:
-        with gr.TabItem("Described Tags"):
-            described_tags = gr.TextArea(label="Described Tags")
-        with gr.TabItem("Original Tags"):
-            original_tags = gr.TextArea(label="Original Tags")
+    transcribe_btn.click(transcribe_and_update, 
+                         inputs=[image_display, image_type, transcriber, booru_tags], 
+                         outputs=[transcribe_image_display, tags_output])
     
-    merged_tags = gr.TextArea(label="Merged Tags")
+    transcribe_with_tags_btn.click(transcribe_and_update, 
+                                   inputs=[image_display, image_type, transcriber, booru_tags], 
+                                   outputs=[transcribe_image_display, tags_output])
     
-    describe_button.click(describe, inputs=[image_input, mode], outputs=[described_tags, original_tags])
-    merge_button.click(merge, inputs=[described_tags, original_tags], outputs=merged_tags)
+    transcribe_btn_final.click(transcribe_image, 
+                               inputs=[transcribe_image_display, image_type, transcriber], 
+                               outputs=[tags_output])
 
-demo.launch()
+app.launch()