app.py

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverSinglestepScheduler
from PIL import Image

# Load a memory-efficient SD variant (under 12GB)
model_id = "runwayml/stable-diffusion-v1-5"

@gr.cache()
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        safety_checker=None,
        use_safetensors=True
    )
    pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cpu")
    pipe.enable_attention_slicing()  # Reduces memory by 30%
    pipe.enable_model_cpu_offload()  # Only loads needed components
    return pipe

def generate_character(prompt, seed=42):
    try:
        pipe = load_model()
        generator = torch.Generator(device="cpu").manual_seed(seed)
        
        with torch.inference_mode():
            image = pipe(
                prompt=f"cartoon character {prompt}, vibrant colors, clean lines",
                negative_prompt="blurry, deformed, ugly",
                num_inference_steps=20,
                guidance_scale=7.5,
                width=512,
                height=512,
                generator=generator
            ).images[0]
        
        return image
    except Exception as e:
        return f"Error: {str(e)}\nTry simplifying your prompt."

# Animation through img2img
def generate_animation(prompt, frames=3):
    base_image = generate_character(prompt)
    if isinstance(base_image, str):  # If error
        return base_image
    
    images = [base_image]
    pipe = load_model()
    
    for i in range(1, frames):
        result = pipe(
            prompt=prompt,
            image=images[-1],
            strength=0.3,  # Small changes per frame
            generator=torch.Generator().manual_seed(i)
        )
        images.append(result.images[0])
    
    images[0].save(
        "animation.gif",
        save_all=True,
        append_images=images[1:],
        duration=500,
        loop=0
    )
    return "animation.gif"

with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Markdown("# 🎬 Character Animator (12GB Optimized)")
    
    with gr.Row():
        prompt = gr.Textbox(
            label="Character Description",
            placeholder="e.g. 'cyberpunk fox wearing sunglasses'"
        )
    
    with gr.Tab("Single Image"):
        img_out = gr.Image(label="Generated Character", type="pil")
        gen_btn = gr.Button("Generate")
    
    with gr.Tab("Animation"):
        anim_out = gr.Image(label="Animation", format="gif")
        anim_btn = gr.Button("Create Animation (3 frames)")
    
    gen_btn.click(generate_character, inputs=prompt, outputs=img_out)
    anim_btn.click(generate_animation, inputs=prompt, outputs=anim_out)

demo.launch()