Update gradio-interface.py

gradio-interface.py

@@ -1,42 +1,116 @@
-import torch
+import os
 import gradio as gr
-from transformers import AutoTokenizer
-from torchvision import transforms
 from PIL import Image
+import torch
+from diffusers import DiffusionPipeline
+import tempfile
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = Animator2D().to(device)
-model.load_state_dict(torch.load("animator2D-model.pth", map_location=device))
-model.eval()
+# Check for GPU availability
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+def initialize_model():
+    """Initialize the Animator2D model."""
+    try:
+        # Initialize the pipeline
+        pipeline = DiffusionPipeline.from_pretrained(
+            "Lod34/Animator2D",
+            trust_remote_code=True,
+            device=DEVICE
+        )
+        return pipeline
+    except Exception as e:
+        raise Exception(f"Error initializing model: {str(e)}")
 
-def generate_sprite(num_frames, description, action, direction):
-    text = f"{num_frames}-frame sprite animation of: {description}, that: {action}, facing: {direction}"
-    encoded_text = tokenizer(
-        text, padding="max_length", max_length=128, truncation=True, return_tensors="pt"
-    )
-    
-    with torch.no_grad():
-        text_ids = encoded_text['input_ids'].to(device)
-        text_mask = encoded_text['attention_mask'].to(device)
-        generated_sprite = model(text_ids, text_mask).cpu().squeeze(0)
-    
-    generated_sprite = (generated_sprite + 1) / 2  # Denormalizzazione
-    generated_sprite = transforms.ToPILImage()(generated_sprite)
-    return generated_sprite
+def generate_animation(
+    description: str,
+    action: str,
+    direction: str,
+    num_frames: int
+):
+    """Generate animation based on input parameters."""
+    try:
+        # Input validation
+        if not all([description, action, direction]):
+            raise ValueError("All text fields must be filled")
+        
+        # Initialize model
+        pipeline = initialize_model()
+        
+        # Prepare prompt
+        prompt = f"A sprite of {description} {action}, facing {direction}"
+        
+        # Generate animation
+        output = pipeline(
+            prompt=prompt,
+            num_frames=num_frames,
+            num_inference_steps=50
+        )
+        
+        # Save animation as GIF
+        temp_dir = tempfile.mkdtemp()
+        output_path = os.path.join(temp_dir, "animation.gif")
+        
+        # Convert output frames to GIF
+        frames = [Image.fromarray(frame) for frame in output.frames]
+        frames[0].save(
+            output_path,
+            save_all=True,
+            append_images=frames[1:],
+            duration=100,
+            loop=0
+        )
+        
+        return output_path
+        
+    except Exception as e:
+        raise gr.Error(f"Generation failed: {str(e)}")
 
-iface = gr.Interface(
-    fn=generate_sprite,
-    inputs=[
-        gr.Number(label="Numero di Frame", value=17),
-        gr.Textbox(label="Descrizione dello Sprite"),
-        gr.Dropdown(["cammina", "corre", "salta", "attacca"], label="Azione"),
-        gr.Dropdown(["Nord", "Sud", "Est", "Ovest"], label="Direzione")
-    ],
-    outputs=gr.Image(type="pil"),
-    title="Animator2D Generator",
-    description="Genera animazioni di sprite basate su descrizioni testuali."
-)
+def create_interface():
+    """Create and launch the Gradio interface."""
+    
+    with gr.Blocks(title="Animator2D Sprite Generator") as interface:
+        gr.Markdown("# Animator2D Sprite Generator")
+        gr.Markdown("Generate animated sprites using AI")
+        
+        with gr.Row():
+            with gr.Column():
+                # Input components
+                description = gr.Textbox(
+                    label="Sprite Description",
+                    placeholder="E.g., a cute pixel art cat"
+                )
+                action = gr.Textbox(
+                    label="Sprite Action",
+                    placeholder="E.g., walking, jumping"
+                )
+                direction = gr.Dropdown(
+                    label="Direction",
+                    choices=["North", "South", "East", "West"],
+                    value="South"
+                )
+                num_frames = gr.Slider(
+                    label="Number of Frames",
+                    minimum=2,
+                    maximum=24,
+                    value=8,
+                    step=1
+                )
+                generate_btn = gr.Button("Generate Animation")
+            
+            with gr.Column():
+                # Output components
+                output_image = gr.Image(label="Generated Animation", type="filepath")
+        
+        # Connect components
+        generate_btn.click(
+            fn=generate_animation,
+            inputs=[description, action, direction, num_frames],
+            outputs=output_image
+        )
+        
+    return interface
 
-iface.launch(share=False)  # Disabilita la condivisione pubblica
+# Launch the application
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch(share=True)