Create app.py

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoProcessor
+import torch
+from PIL import Image
+import io
+
+# Load model and processor (using CPU)
+folder_path = "diffusers/shot-categorizer-v0"
+model = AutoModelForCausalLM.from_pretrained(folder_path, trust_remote_code=True).eval()
+processor = AutoProcessor.from_pretrained(folder_path, trust_remote_code=True)
+
+# Define analysis function
+def analyze_image(image):
+    # Convert Gradio image input to PIL Image
+    if isinstance(image, Image.Image):
+        img = image.convert("RGB")
+    else:
+        img = Image.open(io.BytesIO(image)).convert("RGB")
+
+    prompts = ["<COLOR>", "<LIGHTING>", "<LIGHTING_TYPE>", "<COMPOSITION>"]
+    results = {}
+    
+    # Process each prompt
+    with torch.no_grad():
+        for prompt in prompts:
+            inputs = processor(text=prompt, images=img, return_tensors="pt")
+            generated_ids = model.generate(
+                input_ids=inputs["input_ids"],
+                pixel_values=inputs["pixel_values"],
+                max_new_tokens=1024,
+                early_stopping=False,
+                do_sample=False,
+                num_beams=3,
+            )
+            generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+            parsed_answer = processor.post_process_generation(
+                generated_text, task=prompt, image_size=(img.width, img.height)
+            )
+            results[prompt] = parsed_answer
+
+    # Format the output
+    output_text = "Image Analysis Results:\n\n"
+    output_text += f"Color: {results['<COLOR>']}\n"
+    output_text += f"Lighting: {results['<LIGHTING>']}\n"
+    output_text += f"Lighting Type: {results['<LIGHTING_TYPE>']}\n"
+    output_text += f"Composition: {results['<COMPOSITION>']}\n"
+    
+    return output_text
+
+# Create Gradio interface
+with gr.Blocks(title="Image Analyzer") as demo:
+    gr.Markdown("# Image Analysis Demo")
+    gr.Markdown("Upload an image to analyze its color, lighting, and composition characteristics.")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Image")
+            analyze_button = gr.Button("Analyze Image")
+        
+        with gr.Column():
+            output_text = gr.Textbox(label="Analysis Results", lines=10)
+    
+    # Add example images
+    examples = gr.Examples(
+        examples=["./assets/image_3.jpg"],
+        inputs=image_input,
+        label="Try with this example"
+    )
+    
+    # Connect the button to the function
+    analyze_button.click(
+        fn=analyze_image,
+        inputs=image_input,
+        outputs=output_text
+    )
+
+# Launch the demo
+demo.launch()