import sys import os from utils.wrapper import StreamV2VWrapper import torch from config import Args from pydantic import BaseModel, Field from PIL import Image import math base_model = "runwayml/stable-diffusion-v1-5" default_prompt = "A man is talking" page_content = """

StreamV2V by Jeff-LiangF

Duplicate this space for fast and private usage - thank you!

This demo showcases StreamV2V video-to-video pipeline using 4-step LCM LORA with a MJPEG stream server.

The base model is SD 1.5. We also build in some LORAs for better stylization.

""" class Pipeline: class Info(BaseModel): name: str = "StreamV2V" input_mode: str = "image" page_content: str = page_content class InputParams(BaseModel): prompt: str = Field( default_prompt, title="Prompt", field="textarea", id="prompt", ) # negative_prompt: str = Field( # default_negative_prompt, # title="Negative Prompt", # field="textarea", # id="negative_prompt", # ) width: int = Field( 512, min=2, max=15, title="Width", disabled=True, hide=True, id="width" ) height: int = Field( 512, min=2, max=15, title="Height", disabled=True, hide=True, id="height" ) def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype): params = self.InputParams() self.stream = StreamV2VWrapper( model_id_or_path=base_model, t_index_list=[30, 35, 40, 45], frame_buffer_size=1, width=params.width, height=params.height, warmup=10, acceleration=args.acceleration, do_add_noise=True, output_type="pil", use_denoising_batch=True, use_cached_attn=True, use_feature_injection=True, feature_injection_strength=0.8, feature_similarity_threshold=0.98, cache_interval=4, cache_maxframes=1, use_tome_cache=True, seed=1, ) self._init_lora() self.last_prompt = default_prompt self.stream.prepare( prompt=default_prompt, num_inference_steps=50, guidance_scale=1.0, ) self.lora_active = False self.lora_trigger_words = ['pixelart', 'pixel art', 'Pixel art', 'PixArFK' 'lowpoly', 'low poly', 'Low poly', 'Claymation', 'claymation', 'crayons', 'Crayons', 'crayons doodle', 'Crayons doodle', 'sketch', 'Sketch', 'pencil drawing', 'Pencil drawing', 'oil painting', 'Oil painting'] def _init_lora(self): self.stream.stream.load_lora("./lora_weights/PixelArtRedmond15V-PixelArt-PIXARFK.safetensors", adapter_name='pixelart') self.stream.stream.load_lora("./lora_weights/low_poly.safetensors", adapter_name='lowpoly') self.stream.stream.load_lora("./lora_weights/Claymation.safetensors", adapter_name='claymation') self.stream.stream.load_lora("./lora_weights/doodle.safetensors", adapter_name='crayons') self.stream.stream.load_lora("./lora_weights/Sketch_offcolor.safetensors", adapter_name='sketch') self.stream.stream.load_lora("./lora_weights/bichu-v0612.safetensors", adapter_name='oilpainting') def _activate_lora(self, prompt: str): if any(word in prompt for word in ['pixelart', 'pixel art', 'Pixel art', 'PixArFK']): self.stream.stream.pipe.set_adapters(["lcm", "pixelart"], adapter_weights=[1.0, 1.0]) print("Use LORA: pixelart in ./lora_weights/PixelArtRedmond15V-PixelArt-PIXARFK.safetensors") elif any(word in prompt for word in ['lowpoly', 'low poly', 'Low poly']): self.stream.stream.pipe.set_adapters(["lcm", "lowpoly"], adapter_weights=[1.0, 1.0]) print("Use LORA: lowpoly in ./lora_weights/low_poly.safetensors") elif any(word in prompt for word in ['Claymation', 'claymation']): self.stream.stream.pipe.set_adapters(["lcm", "claymation"], adapter_weights=[1.0, 1.0]) print("Use LORA: claymation in ./lora_weights/Claymation.safetensors") elif any(word in prompt for word in ['crayons', 'Crayons', 'crayons doodle', 'Crayons doodle']): self.stream.stream.pipe.set_adapters(["lcm", "crayons"], adapter_weights=[1.0, 1.0]) print("Use LORA: crayons in ./lora_weights/doodle.safetensors") elif any(word in prompt for word in ['sketch', 'Sketch', 'pencil drawing', 'Pencil drawing']): self.stream.stream.pipe.set_adapters(["lcm", "sketch"], adapter_weights=[1.0, 1.0]) print("Use LORA: sketch in ./lora_weights/Sketch_offcolor.safetensors") elif any(word in prompt for word in ['oil painting', 'Oil painting']): self.stream.stream.pipe.set_adapters(["lcm", "oilpainting"], adapter_weights=[1.0, 1.0]) print("Use LORA: oilpainting in ./lora_weights/bichu-v0612.safetensors") def _deactivate_lora(self): self.stream.stream.pipe.set_adapters("lcm") print("Deactivate LORA, back to SD1.5") def _check_trigger_words(self, prompt: str): return any(word in prompt for word in self.lora_trigger_words) def predict(self, params: "Pipeline.InputParams") -> Image.Image: if self._check_trigger_words(params.prompt): if not self.lora_active: self._activate_lora(params.prompt) self.lora_active = True else: if self.lora_active: self._deactivate_lora() self.lora_active = False image_tensor = self.stream.preprocess_image(params.image) output_image = self.stream(image=image_tensor, prompt=params.prompt) return output_image