idisc/Multi-View-Illusion-Diffusion

1from diffusers import DiffusionPipeline2from diffusers.utils import pt_to_pil3import gradio as gr4import torch5import numpy as np6 7 8stage_1 = DiffusionPipeline.from_pretrained(9    "DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float1610)11stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.012stage_1.enable_model_cpu_offload()13stage_2 = DiffusionPipeline.from_pretrained(14    "DeepFloyd/IF-II-M-v1.0",15    text_encoder=None,16    variant="fp16",17    torch_dtype=torch.float16,18)19stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.020stage_2.enable_model_cpu_offload()21 22# stage 323safety_modules = {24    "feature_extractor": stage_1.feature_extractor,25    "safety_checker": stage_1.safety_checker,26    "watermarker": stage_1.watermarker,27}28stage_3 = DiffusionPipeline.from_pretrained(29    "stabilityai/stable-diffusion-x4-upscaler",30    **safety_modules,31    torch_dtype=torch.float1632)33stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.034stage_3.enable_model_cpu_offload()35 36 37def predict(prompt):38    prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)39    generator = torch.manual_seed(0)40    image = stage_1(41        prompt_embeds=prompt_embeds,42        negative_prompt_embeds=negative_embeds,43        generator=generator,44        output_type="pt",45    ).images46    image = stage_2(47        image=image,48        prompt_embeds=prompt_embeds,49        negative_prompt_embeds=negative_embeds,50        generator=generator,51        output_type="pt",52    ).images53    image = stage_3(54        prompt=prompt, image=image, generator=generator, noise_level=10055    ).images[0]56    return image57 58 59gradio_app = gr.Interface(60    fn=predict,61    inputs="text",62    outputs="image",63    title="Text to Image Generator",64    description="Enter a text string to generate an image.",65)66 67if __name__ == "__main__":68    gradio_app.launch(server_name="0.0.0.0") # server_name="0.0.0.0"69