mung-bean/sceneweaver

1from diffusers import (2    StableDiffusionXLPipeline,3    StableDiffusionXLAdapterPipeline,4    AutoencoderKL,5    UniPCMultistepScheduler,6    T2IAdapter,7)8import torch, os9from PIL import Image10from io import BytesIO11import models12from database import SessionLocal13from text_processor import (14    get_resolved_sentences,15    detect_and_translate_to_english,16    get_script_captions,17)18from s3 import upload_image_to_s319from diffusers.utils import load_image20import random21from controlnet_aux import OpenposeDetector22import numpy as np23import gc24 25# Global device configuration26 27dtype = torch.float1628 29# Initialize global generator30generator = torch.Generator()31 32# Initialize the models globally to ensure they're only loaded once33print("Loading VAE...")34vae = AutoencoderKL.from_pretrained(35    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=dtype, use_safetensors=True36).to("cuda")37 38print("Loading base pipeline...")39pipe = StableDiffusionXLPipeline.from_pretrained(40    "stabilityai/stable-diffusion-xl-base-1.0",41    vae=vae,42    torch_dtype=dtype,43    variant="fp16",44    use_safetensors=True,45).to("cuda")46 47pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)48 49 50pipe.load_lora_weights("safetensors/Storyboard_sketch.safetensors", adapter_name="sketch")51pipe.load_lora_weights("safetensors/anglesv2.safetensors", adapter_name="angles")52pipe.set_adapters(["sketch", "angles"], adapter_weights=[0.5, 0.5])53pipe.enable_xformers_memory_efficient_attention()54 55print("Loading OpenPose detector...")56openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")57 58print("Loading T2I adapter...")59adapter = T2IAdapter.from_pretrained(60    "TencentARC/t2i-adapter-openpose-sdxl-1.0", torch_dtype=dtype61).to("cuda")62 63print("Loading adapter pipeline...")64posepipe = StableDiffusionXLAdapterPipeline.from_pretrained(65    "stabilityai/stable-diffusion-xl-base-1.0",66    adapter=adapter,67    vae=vae,68    torch_dtype=dtype,69    variant="fp16",70    use_safetensors=True,71).to("cuda")72 73 74posepipe.scheduler = UniPCMultistepScheduler.from_config(posepipe.scheduler.config)75 76posepipe.load_lora_weights(77    "safetensors/Storyboard_sketch.safetensors", adapter_name="sketch"78)79posepipe.load_lora_weights("safetensors/anglesv2.safetensors", adapter_name="angles")80posepipe.set_adapters(["sketch", "angles"], adapter_weights=[0.5, 0.5])81posepipe.enable_xformers_memory_efficient_attention()82 83print("All models loaded successfully")84 85 86def clear_cuda_cache():87    """Clear CUDA cache to free up memory"""88    if torch.cuda.is_available():89        torch.cuda.empty_cache()90        gc.collect()91 92 93def get_dimensions(resolution: str) -> tuple[int, int]:94    resolution_map = {95        "16:9": (1024, 576),96        "1:1": (1024, 1024),97        "9:16": (576, 1024),98    }99    return resolution_map.get(resolution, (1024, 1024))100 101 102def generate_batch_images(103    story: str, storyboard_id: int, resolution: str = "1:1", isStory: bool = True104):105    # Clear cache before batch generation106    clear_cuda_cache()107 108    db = SessionLocal()109    try:110        if isStory:111            prompts = get_resolved_sentences(story)112        elif not isStory:113            prompts = get_script_captions(story)114 115        width, height = get_dimensions(resolution)116 117        for num, prompt in enumerate(prompts):118            # Generate a random seed for each image in the batch119            seed = random.randint(0, 2**32 - 1)120            generator.manual_seed(seed)121 122            print(f"Generating image {num+1} with seed {seed}")123 124            result = pipe(125                prompt=f"Storyboard sketch of {prompt}, black and white, cinematic, high quality",126                negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy, abstract, bad physics",127                guidance_scale=8.5,128                height=height,129                width=width,130                num_inference_steps=30,131                generator=generator,132            )133 134            image = result.images[0]135            buf = BytesIO()136            image.save(buf, format="JPEG")137            buf.seek(0)138 139            s3_url = upload_image_to_s3(140                buf.read(),141                f"image_{num + 1}.jpg",142                folder=f"storyboards/{storyboard_id}",143            )144 145            db_image = models.Image(146                storyboard_id=storyboard_id,147                image_path=s3_url,148                caption=prompt,149            )150            db.add(db_image)151            db.commit()152            db.refresh(db_image)153 154            print(f"Image {num+1} generated successfully")155 156            # Clear cache after each image157            clear_cuda_cache()158 159    except Exception as e:160        print(f"Error during image generation: {e}")161        import traceback162 163        traceback.print_exc()164        db.rollback()165    finally:166        db.close()167 168 169def generate_single_image(170    image_id: int,171    caption: str,172    seed: int = None,173    resolution: str = "1:1",174    isOpenPose: bool = False,175    pose_img: Image.Image = None,176):177    # Clear cache before single image generation178    clear_cuda_cache()179 180    db = SessionLocal()181    try:182        # Get existing image record183        db_image = db.query(models.Image).filter(models.Image.id == image_id).first()184        processed_caption = detect_and_translate_to_english(caption)185        width, height = get_dimensions(resolution)186 187        # Use provided seed or generate a random one188        current_seed = seed if seed is not None else random.randint(0, 2**32 - 1)189        generator.manual_seed(current_seed)190 191        print(f"Generating single image with seed {current_seed}")192 193        if not db_image:194            raise ValueError(f"Image with id {image_id} not found.")195 196        if isOpenPose:197            print("Using OpenPose pipeline")198            image = openpose(pose_img, detect_resolution=512, image_resolution=1024)199            image = np.array(image)[:, :, ::-1]200            image = Image.fromarray(np.uint8(image))201 202            result = posepipe(203                prompt=f"Storyboard sketch of {processed_caption}, black and white, cinematic, high quality",204                negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy, abstract, bad physics",205                image=image,206                adapter_conditioning_scale=1,207                guidance_scale=8.5,208                num_inference_steps=30,209                generator=generator,210            )211        else:212            print("Using standard pipeline")213            result = pipe(214                prompt=f"Storyboard sketch of {processed_caption}, black and white, cinematic, high quality",215                negative_prompt="ugly, deformed, disfigured, poor details, bad anatomy, abstract, bad physics",216                guidance_scale=8.5,217                num_inference_steps=30,218                width=width,219                height=height,220                generator=generator,221            )222 223        # Save and upload224        image = result.images[0]225        buf = BytesIO()226        image.save(buf, format="JPEG")227        buf.seek(0)228 229        s3_url = upload_image_to_s3(230            buf.read(),231            f"image_{image_id}.jpg",232            folder=f"storyboards/{db_image.storyboard_id}",233        )234 235        # Update image record236        db_image.image_path = s3_url237        db_image.caption = caption238        db_image.seed = current_seed239        db.commit()240        db.refresh(db_image)241 242        print(f"Single image generated successfully")243 244        # Clear cache after generation245        clear_cuda_cache()246 247        return db_image248 249    except Exception as e:250        print(f"Error during image regeneration: {e}")251        import traceback252 253        traceback.print_exc()254        db.rollback()255        return None256    finally:257        db.close()258